fix(grok): 视频在 status 返回 video.url 后再按上游参数计费

创建/generations 仅绑定账号并缓存 create-time 定价快照;轮询到
video.url / download_url 等成功 URL 时一次性扣费。时长/分辨率/模型
优先取上游 status 字段,缺失则回退创建快照,再套默认 8s/480p。
Redis SetNX + request_id 去重,避免重复轮询双扣。
This commit is contained in:
IanShaw027
2026-08-07 16:29:57 +08:00
parent 25d2b03e90
commit b59a3df5b5
15 changed files with 601 additions and 22 deletions
+121 -9
View File
@@ -406,7 +406,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, result), true, nil)
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
if isGrokVideoCreateEndpoint(endpoint) && strings.TrimSpace(result.ResponseID) != "" {
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(
requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID,
); err != nil {
@@ -416,8 +416,28 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
zap.Error(err),
)
}
// Defer billing until status polling observes video.url. Persist create-time
// model/duration/resolution so status can still price if upstream omits them.
if err := h.gatewayService.StoreGrokVideoPendingBilling(requestCtx, result.ResponseID, subject.UserID, apiKey.ID, service.GrokVideoPendingBilling{
Model: requestModel,
BillingModel: firstNonEmptyString(result.BillingModel, requestModel),
UpstreamModel: result.UpstreamModel,
VideoResolution: result.VideoResolution,
VideoDurationSeconds: result.VideoDurationSeconds,
OriginalModel: clientRequestedModel(c, requestModel),
}); err != nil {
reqLog.Warn("grok_media.store_video_pending_billing_failed",
zap.Int64("account_id", account.ID),
zap.String("request_id", result.ResponseID),
zap.Error(err),
)
}
}
if shouldRecordGrokMediaUsage(endpoint, requestModel, result) {
if endpoint == service.GrokMediaEndpointVideoStatus {
if billResult := prepareGrokVideoStatusBilling(requestCtx, h, reqLog, apiKey, subject, requestID, result); billResult != nil {
recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, billResult, billResult.Model, body, requestID)
}
} else if shouldRecordGrokMediaUsage(endpoint, requestModel, result) {
recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, result, requestModel, body, requestID)
}
reqLog.Debug("grok_media.request_completed",
@@ -459,22 +479,114 @@ func grokMediaScheduleModel(account *service.Account, routingModel string, resul
return account.GetMappedModel(routingModel)
}
// shouldRecordGrokMediaUsage gates usage writes for Grok media generation only.
// Status/content polls, empty model, and failed generations with zero billable
// image/video units never bill.
func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string, result *service.OpenAIForwardResult) bool {
if !endpoint.IsGenerationRequest() || strings.TrimSpace(requestModel) == "" {
func isGrokVideoCreateEndpoint(endpoint service.GrokMediaEndpoint) bool {
switch endpoint {
case service.GrokMediaEndpointVideosGenerations,
service.GrokMediaEndpointVideosEdits,
service.GrokMediaEndpointVideosExtensions:
return true
default:
return false
}
}
// shouldRecordGrokMediaUsage gates usage writes for immediate (image) generation.
// Async video create never bills here — status polling does when video.url appears.
// Status/content polls, empty model, and failed generations with zero billable
// image units never bill via this helper.
func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string, result *service.OpenAIForwardResult) bool {
if result == nil {
return false
}
if result.VideoCount > 0 {
return true
if isGrokVideoCreateEndpoint(endpoint) || endpoint.IsVideoLookupRequest() {
return false
}
if !endpoint.IsGenerationRequest() || strings.TrimSpace(requestModel) == "" {
return false
}
return result.ImageCount > 0
}
// prepareGrokVideoStatusBilling claims one-shot billing and merges status-body
// duration/resolution/model with the create-time pending snapshot.
func prepareGrokVideoStatusBilling(
ctx context.Context,
h *OpenAIGatewayHandler,
reqLog *zap.Logger,
apiKey *service.APIKey,
subject middleware2.AuthSubject,
requestID string,
statusResult *service.OpenAIForwardResult,
) *service.OpenAIForwardResult {
if h == nil || h.gatewayService == nil || apiKey == nil || statusResult == nil {
return nil
}
// Prefer units already extracted on the forward result; otherwise re-check pending-only path
// is not enough without a status URL (VideoCount stays 0).
if statusResult.VideoCount <= 0 {
return nil
}
claimed, err := h.gatewayService.ClaimGrokVideoBilling(ctx, requestID, subject.UserID, apiKey.ID)
if err != nil {
reqLog.Warn("grok_media.video_billing_claim_failed", zap.String("request_id", requestID), zap.Error(err))
return nil
}
if !claimed {
reqLog.Debug("grok_media.video_billing_already_claimed", zap.String("request_id", requestID))
return nil
}
pending, loadErr := h.gatewayService.LoadGrokVideoPendingBilling(ctx, requestID, subject.UserID, apiKey.ID)
if loadErr != nil {
reqLog.Warn("grok_media.video_pending_billing_load_failed", zap.String("request_id", requestID), zap.Error(loadErr))
}
// Re-merge with pending so create-time model wins when status omits model fields.
// statusResult already has status-body duration/resolution when present.
merged := *statusResult
if pending != nil {
if strings.TrimSpace(merged.Model) == "" {
merged.Model = firstNonEmptyString(pending.BillingModel, pending.Model, pending.OriginalModel)
}
if strings.TrimSpace(merged.BillingModel) == "" {
merged.BillingModel = firstNonEmptyString(pending.BillingModel, pending.Model, merged.Model)
}
if strings.TrimSpace(merged.UpstreamModel) == "" {
merged.UpstreamModel = pending.UpstreamModel
}
if strings.TrimSpace(merged.VideoResolution) == "" {
merged.VideoResolution = pending.VideoResolution
}
if merged.VideoDurationSeconds <= 0 {
merged.VideoDurationSeconds = pending.VideoDurationSeconds
}
if strings.TrimSpace(merged.ResponseID) == "" {
merged.ResponseID = requestID
}
}
if strings.TrimSpace(merged.Model) == "" {
merged.Model = "grok-imagine-video"
}
if strings.TrimSpace(merged.BillingModel) == "" {
merged.BillingModel = merged.Model
}
// Stable dedup key: same video task bills once across status polls.
if strings.TrimSpace(merged.RequestID) == "" {
merged.RequestID = "grok-video:" + strings.TrimSpace(firstNonEmptyString(merged.ResponseID, requestID))
}
merged.VideoCount = 1
merged.VideoResolution = service.NormalizeVideoBillingResolutionOrDefault(merged.VideoResolution)
merged.VideoDurationSeconds = service.NormalizeVideoBillingDurationSecondsOrDefault(merged.VideoDurationSeconds)
return &merged
}
func firstNonEmptyString(values ...string) string {
for _, v := range values {
if s := strings.TrimSpace(v); s != "" {
return s
}
}
return ""
}
func recordGrokMediaUsage(
c *gin.Context,
h *OpenAIGatewayHandler,
+5 -5
View File
@@ -42,13 +42,13 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
want: true,
},
{
name: "video generation records usage",
name: "video generation defers usage until status",
endpoint: service.GrokMediaEndpointVideosGenerations,
model: "grok-imagine-video-1.5",
want: true,
want: false,
},
{
name: "video status skips empty model usage",
name: "video status skips immediate helper (status path claims separately)",
endpoint: service.GrokMediaEndpointVideoStatus,
model: "",
want: false,
@@ -71,9 +71,9 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
t.Run(tt.name, func(t *testing.T) {
// Nil result must never bill.
require.False(t, shouldRecordGrokMediaUsage(tt.endpoint, tt.model, nil))
// Non-nil result with units only bills generation endpoints with model.
// Immediate helper only bills image generation (async video bills on status).
result := &service.OpenAIForwardResult{ImageCount: 1, VideoCount: 0}
if tt.endpoint.IsGenerationRequest() && strings.TrimSpace(tt.model) != "" {
if tt.endpoint.IsGenerationRequest() && !isGrokVideoCreateEndpoint(tt.endpoint) && strings.TrimSpace(tt.model) != "" {
require.Equal(t, tt.want, shouldRecordGrokMediaUsage(tt.endpoint, tt.model, result))
} else {
require.False(t, shouldRecordGrokMediaUsage(tt.endpoint, tt.model, result))
@@ -7,6 +7,7 @@ import (
"errors"
"fmt"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -64,6 +65,57 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
return c.rdb.Del(ctx, key).Err()
}
const (
grokVideoPendingBillingPrefix = "grok_video_pending:"
grokVideoBilledPrefix = "grok_video_billed:"
)
func (c *gatewayCache) SetGrokVideoPendingBilling(ctx context.Context, key string, payload []byte, ttl time.Duration) error {
if c == nil || c.rdb == nil {
return errors.New("gateway cache unavailable")
}
key = strings.TrimSpace(key)
if key == "" || len(payload) == 0 {
return errors.New("invalid grok video pending billing payload")
}
if ttl <= 0 {
ttl = 24 * time.Hour
}
return c.rdb.Set(ctx, grokVideoPendingBillingPrefix+key, payload, ttl).Err()
}
func (c *gatewayCache) GetGrokVideoPendingBilling(ctx context.Context, key string) ([]byte, error) {
if c == nil || c.rdb == nil {
return nil, errors.New("gateway cache unavailable")
}
key = strings.TrimSpace(key)
if key == "" {
return nil, errors.New("invalid grok video pending billing key")
}
val, err := c.rdb.Get(ctx, grokVideoPendingBillingPrefix+key).Bytes()
if err != nil {
if errors.Is(err, redis.Nil) {
return nil, nil
}
return nil, err
}
return val, nil
}
func (c *gatewayCache) ClaimGrokVideoBilled(ctx context.Context, key string, ttl time.Duration) (bool, error) {
if c == nil || c.rdb == nil {
return false, errors.New("gateway cache unavailable")
}
key = strings.TrimSpace(key)
if key == "" {
return false, errors.New("invalid grok video billed key")
}
if ttl <= 0 {
ttl = 48 * time.Hour
}
return c.rdb.SetNX(ctx, grokVideoBilledPrefix+key, "1", ttl).Result()
}
// Compile-time assertion: gatewayCache must implement CyberSessionBlockStore.
var _ service.CyberSessionBlockStore = (*gatewayCache)(nil)
var _ service.LiveCallStore = (*gatewayCache)(nil)
@@ -144,6 +144,17 @@ func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Conte
return nil
}
func (s *stickyGatewayCacheHotpathStub) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (s *stickyGatewayCacheHotpathStub) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (s *stickyGatewayCacheHotpathStub) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) {
s.listByGroupCalls.Add(1)
if s.err != nil {
@@ -277,6 +277,16 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context
return nil
}
func (m *mockGatewayCacheForPlatform) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (m *mockGatewayCacheForPlatform) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (m *mockGatewayCacheForPlatform) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
type mockGroupRepoForGateway struct {
groups map[int64]*Group
getByIDCalls int
@@ -476,6 +476,15 @@ type GatewayCache interface {
// DeleteSessionAccountID 删除粘性会话绑定,用于账号不可用时主动清理
// Delete sticky session binding, used to proactively clean up when account becomes unavailable
DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error
// Grok async video billing snapshot (create → status success).
// SetGrokVideoPendingBilling stores create-time model/duration/resolution for status billing.
SetGrokVideoPendingBilling(ctx context.Context, key string, payload []byte, ttl time.Duration) error
// GetGrokVideoPendingBilling returns the create-time billing snapshot; miss → nil, nil.
GetGrokVideoPendingBilling(ctx context.Context, key string) ([]byte, error)
// ClaimGrokVideoBilled atomically marks a video request as billed (SetNX).
// Returns true when this caller won the claim; false when already billed or claim unavailable.
ClaimGrokVideoBilled(ctx context.Context, key string, ttl time.Duration) (bool, error)
}
// derefGroupID safely dereferences *int64 to int64, returning 0 if nil
@@ -305,6 +305,16 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context,
return nil
}
func (m *mockGatewayCacheForGemini) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (m *mockGatewayCacheForGemini) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (m *mockGatewayCacheForGemini) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
// TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择
func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) {
ctx := context.Background()
+241 -5
View File
@@ -14,6 +14,7 @@ import (
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
@@ -315,6 +316,215 @@ func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount(
return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey)
}
// GrokVideoPendingBilling is the create-time snapshot used when status polling
// first observes a completed video URL. Status may omit model/duration; we fall
// back to this snapshot, then defaults.
type GrokVideoPendingBilling struct {
Model string `json:"model"`
BillingModel string `json:"billing_model,omitempty"`
UpstreamModel string `json:"upstream_model,omitempty"`
VideoResolution string `json:"video_resolution,omitempty"`
VideoDurationSeconds int `json:"video_duration_seconds,omitempty"`
OriginalModel string `json:"original_model,omitempty"`
}
func grokVideoPendingBillingKey(requestID string, userID, apiKeyID int64) string {
requestID = strings.TrimSpace(requestID)
if requestID == "" || userID <= 0 || apiKeyID <= 0 {
return ""
}
return fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID)
}
func grokVideoPendingBillingTTL(cfg *config.Config) time.Duration {
// Video generation can take several minutes; keep create-time pricing for a day.
_ = cfg
return 24 * time.Hour
}
func grokVideoBilledClaimTTL(cfg *config.Config) time.Duration {
_ = cfg
return 48 * time.Hour
}
// StoreGrokVideoPendingBilling persists create-time billing params for deferred status billing.
func (s *OpenAIGatewayService) StoreGrokVideoPendingBilling(
ctx context.Context,
requestID string,
userID, apiKeyID int64,
pending GrokVideoPendingBilling,
) error {
if s == nil || s.cache == nil {
return fmt.Errorf("grok video pending billing cache is unavailable")
}
key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
if key == "" {
return fmt.Errorf("grok video pending billing key is invalid")
}
pending.Model = strings.TrimSpace(pending.Model)
pending.BillingModel = strings.TrimSpace(pending.BillingModel)
pending.UpstreamModel = strings.TrimSpace(pending.UpstreamModel)
pending.OriginalModel = strings.TrimSpace(pending.OriginalModel)
if pending.VideoResolution != "" {
pending.VideoResolution = NormalizeVideoBillingResolutionOrDefault(pending.VideoResolution)
}
if pending.VideoDurationSeconds > 0 {
pending.VideoDurationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(pending.VideoDurationSeconds)
}
payload, err := json.Marshal(pending)
if err != nil {
return err
}
return s.cache.SetGrokVideoPendingBilling(ctx, key, payload, grokVideoPendingBillingTTL(s.cfg))
}
// LoadGrokVideoPendingBilling returns the create-time snapshot (may be nil on miss).
func (s *OpenAIGatewayService) LoadGrokVideoPendingBilling(
ctx context.Context,
requestID string,
userID, apiKeyID int64,
) (*GrokVideoPendingBilling, error) {
if s == nil || s.cache == nil {
return nil, fmt.Errorf("grok video pending billing cache is unavailable")
}
key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
if key == "" {
return nil, fmt.Errorf("grok video pending billing key is invalid")
}
payload, err := s.cache.GetGrokVideoPendingBilling(ctx, key)
if err != nil || len(payload) == 0 {
return nil, err
}
var pending GrokVideoPendingBilling
if err := json.Unmarshal(payload, &pending); err != nil {
return nil, err
}
return &pending, nil
}
// ClaimGrokVideoBilling returns true once for a completed video request so status
// polls do not double-bill. Fail-closed: claim errors are treated as already billed.
func (s *OpenAIGatewayService) ClaimGrokVideoBilling(
ctx context.Context,
requestID string,
userID, apiKeyID int64,
) (bool, error) {
if s == nil || s.cache == nil {
return false, fmt.Errorf("grok video billing claim cache is unavailable")
}
key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
if key == "" {
return false, fmt.Errorf("grok video billing claim key is invalid")
}
return s.cache.ClaimGrokVideoBilled(ctx, key, grokVideoBilledClaimTTL(s.cfg))
}
// IsGrokVideoStatusBillable reports whether a status response body has a downloadable
// video URL (upstream success). Empty/pending/failed responses are not billable.
func IsGrokVideoStatusBillable(statusBody []byte) bool {
return strings.TrimSpace(extractGrokVideoStatusContentURL(statusBody)) != ""
}
// extractGrokVideoStatusContentURL returns the first usable video URL from a status body.
func extractGrokVideoStatusContentURL(body []byte) string {
if len(body) == 0 || !gjson.ValidBytes(body) {
return ""
}
for _, path := range []string{
"video.url",
"url",
"download_url",
"video_url",
"data.video.url",
"data.url",
"data.download_url",
} {
if u := strings.TrimSpace(gjson.GetBytes(body, path).String()); u != "" {
return u
}
}
return ""
}
// ExtractGrokVideoBillingFromStatusBody builds usage units from an upstream status
// payload. Prefers status fields; falls back to create-time pending snapshot.
func ExtractGrokVideoBillingFromStatusBody(statusBody []byte, pending *GrokVideoPendingBilling, requestID string) *OpenAIForwardResult {
if !IsGrokVideoStatusBillable(statusBody) {
return nil
}
model := ""
billingModel := ""
upstreamModel := ""
resolution := ""
durationSeconds := 0
if gjson.ValidBytes(statusBody) {
model = firstNonEmpty(
strings.TrimSpace(gjson.GetBytes(statusBody, "model").String()),
strings.TrimSpace(gjson.GetBytes(statusBody, "video.model").String()),
strings.TrimSpace(gjson.GetBytes(statusBody, "data.model").String()),
)
resolution = firstNonEmpty(
strings.TrimSpace(gjson.GetBytes(statusBody, "resolution").String()),
strings.TrimSpace(gjson.GetBytes(statusBody, "video.resolution").String()),
strings.TrimSpace(gjson.GetBytes(statusBody, "data.resolution").String()),
)
for _, path := range []string{"duration", "video.duration", "duration_seconds", "video.duration_seconds", "data.duration"} {
if v := gjson.GetBytes(statusBody, path); v.Exists() && v.Type == gjson.Number {
durationSeconds = int(v.Int())
break
}
}
}
if pending != nil {
if model == "" {
model = firstNonEmpty(pending.BillingModel, pending.Model, pending.OriginalModel)
}
if billingModel == "" {
billingModel = firstNonEmpty(pending.BillingModel, pending.Model)
}
if upstreamModel == "" {
upstreamModel = pending.UpstreamModel
}
if resolution == "" {
resolution = pending.VideoResolution
}
if durationSeconds <= 0 {
durationSeconds = pending.VideoDurationSeconds
}
}
if model == "" {
// Still bill once URL exists; pricing resolver uses model family defaults.
model = "grok-imagine-video"
}
if billingModel == "" {
billingModel = model
}
// Leave empty resolution / zero duration for the handler to merge create-time
// pending snapshot before applying defaults.
if resolution != "" {
resolution = NormalizeVideoBillingResolutionOrDefault(resolution)
}
if durationSeconds > 0 {
durationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(durationSeconds)
}
responseID := extractGrokMediaVideoRequestID(statusBody)
if responseID == "" {
responseID = strings.TrimSpace(requestID)
}
return &OpenAIForwardResult{
ResponseID: responseID,
Model: model,
BillingModel: billingModel,
UpstreamModel: upstreamModel,
VideoCount: 1,
VideoResolution: resolution,
VideoDurationSeconds: durationSeconds,
// Keep legacy media-unit counter for existing usage displays.
ImageCount: 1,
}
}
func (s *OpenAIGatewayService) ForwardGrokMedia(
ctx context.Context,
c *gin.Context,
@@ -437,12 +647,23 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
}
writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter)
usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody)
resultModel := requestModel
resultBillingModel := requestModel
if endpoint == GrokMediaEndpointVideoStatus {
// Status has no request body model; use upstream status fields when billable.
if m := strings.TrimSpace(usage.Model); m != "" {
resultModel = m
}
if m := strings.TrimSpace(usage.BillingModel); m != "" {
resultBillingModel = m
}
}
return &OpenAIForwardResult{
RequestID: requestIDHeader,
ResponseID: usage.ResponseID,
Usage: usage.Usage,
Model: requestModel,
BillingModel: requestModel,
Model: resultModel,
BillingModel: resultBillingModel,
UpstreamModel: upstreamModel,
ResponseHeaders: resp.Header.Clone(),
Duration: time.Since(startTime),
@@ -783,6 +1004,8 @@ func NormalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string
type grokMediaUsageMetadata struct {
ResponseID string
Usage OpenAIUsage
Model string
BillingModel string
ImageCount int
ImageSize string
ImageInputSize string
@@ -802,12 +1025,25 @@ func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMedi
meta.ImageInputSize = requestInfo.Size
meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody)
case GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions:
// Async video: capture request_id + create-time pricing params only.
// Billable VideoCount is set later when status polling observes video.url.
meta.ResponseID = extractGrokMediaVideoRequestID(responseBody)
meta.VideoCount = 1
meta.VideoResolution = requestInfo.Resolution
meta.VideoDurationSeconds = requestInfo.DurationSeconds
// Keep the legacy media-unit counter populated for existing usage displays.
meta.ImageCount = 1
case GrokMediaEndpointVideoStatus:
// Prefer status-body URL success + upstream duration/resolution when present.
if IsGrokVideoStatusBillable(responseBody) {
// provisional units; handler merges with pending snapshot before RecordUsage.
if billed := ExtractGrokVideoBillingFromStatusBody(responseBody, nil, ""); billed != nil {
meta.ResponseID = billed.ResponseID
meta.Model = billed.Model
meta.BillingModel = billed.BillingModel
meta.VideoCount = billed.VideoCount
meta.VideoResolution = billed.VideoResolution
meta.VideoDurationSeconds = billed.VideoDurationSeconds
meta.ImageCount = billed.ImageCount
}
}
}
return meta
}
@@ -0,0 +1,86 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestIsGrokVideoStatusBillable(t *testing.T) {
t.Parallel()
require.False(t, IsGrokVideoStatusBillable(nil))
require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"pending"}`)))
require.False(t, IsGrokVideoStatusBillable([]byte(`{"status":"completed"}`)))
require.True(t, IsGrokVideoStatusBillable([]byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/x.mp4"}}`)))
require.True(t, IsGrokVideoStatusBillable([]byte(`{"url":"https://example.com/v.mp4"}`)))
require.True(t, IsGrokVideoStatusBillable([]byte(`{"download_url":"/v1/videos/task/content"}`)))
}
func TestExtractGrokVideoBillingFromStatusBodyPrefersUpstreamParams(t *testing.T) {
t.Parallel()
pending := &GrokVideoPendingBilling{
Model: "pending-model",
BillingModel: "pending-billing",
UpstreamModel: "pending-upstream",
VideoResolution: VideoBillingResolution480P,
VideoDurationSeconds: 8,
}
body := []byte(`{
"status":"done",
"model":"status-model",
"video":{"url":"https://vidgen.x.ai/signed.mp4","duration":12,"resolution":"720p"}
}`)
result := ExtractGrokVideoBillingFromStatusBody(body, pending, "req-1")
require.NotNil(t, result)
require.Equal(t, 1, result.VideoCount)
require.Equal(t, "status-model", result.Model)
require.Equal(t, VideoBillingResolution720P, result.VideoResolution)
require.Equal(t, 12, result.VideoDurationSeconds)
}
func TestExtractGrokVideoBillingFromStatusBodyFallsBackToPending(t *testing.T) {
t.Parallel()
pending := &GrokVideoPendingBilling{
Model: "create-model",
BillingModel: "create-billing",
UpstreamModel: "create-upstream",
VideoResolution: VideoBillingResolution1080P,
VideoDurationSeconds: 10,
}
body := []byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/signed.mp4"}}`)
result := ExtractGrokVideoBillingFromStatusBody(body, pending, "req-2")
require.NotNil(t, result)
require.Equal(t, "create-billing", result.BillingModel)
require.Equal(t, "create-upstream", result.UpstreamModel)
require.Equal(t, VideoBillingResolution1080P, result.VideoResolution)
require.Equal(t, 10, result.VideoDurationSeconds)
}
func TestGrokMediaUsageFromResponseVideoCreateDoesNotBill(t *testing.T) {
t.Parallel()
info := GrokMediaRequestInfo{Model: "grok-imagine-video", Resolution: "720p", DurationSeconds: 10}
meta := grokMediaUsageFromResponse(GrokMediaEndpointVideosGenerations, info, []byte(`{"request_id":"v1"}`))
require.Equal(t, "v1", meta.ResponseID)
require.Equal(t, 0, meta.VideoCount)
require.Equal(t, 10, meta.VideoDurationSeconds)
require.Equal(t, VideoBillingResolution720P, meta.VideoResolution)
}
func TestGrokMediaUsageFromResponseVideoStatusBillsOnURL(t *testing.T) {
t.Parallel()
meta := grokMediaUsageFromResponse(
GrokMediaEndpointVideoStatus,
GrokMediaRequestInfo{},
[]byte(`{"status":"completed","video":{"url":"https://vidgen.x.ai/a.mp4","duration":9,"resolution":"480p"}}`),
)
require.Equal(t, 1, meta.VideoCount)
require.Equal(t, 9, meta.VideoDurationSeconds)
require.Equal(t, VideoBillingResolution480P, meta.VideoResolution)
pendingOnly := grokMediaUsageFromResponse(
GrokMediaEndpointVideoStatus,
GrokMediaRequestInfo{},
[]byte(`{"status":"completed"}`),
)
require.Equal(t, 0, pendingOnly.VideoCount)
}
@@ -182,6 +182,16 @@ func (c *schedulerTestGatewayCache) DeleteSessionAccountID(ctx context.Context,
return nil
}
func (c *schedulerTestGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c *schedulerTestGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c *schedulerTestGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func newSchedulerTestOpenAIWSV2Config() *config.Config {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.Enabled = true
@@ -133,6 +133,17 @@ func (c *comboCacheAndStore) RefreshSessionTTL(_ context.Context, _ int64, _ str
func (c *comboCacheAndStore) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error {
return nil
}
func (c *comboCacheAndStore) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c *comboCacheAndStore) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c *comboCacheAndStore) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func (c *comboCacheAndStore) SetCyberSessionBlocked(ctx context.Context, key string, ttl time.Duration) error {
return c.store.SetCyberSessionBlocked(ctx, key, ttl)
}
@@ -1209,9 +1209,10 @@ func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T)
require.Equal(t, "grok-imagine-video", result.BillingModel)
require.Equal(t, 3, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.OutputTokens)
require.Equal(t, 1, result.ImageCount)
// Create accepts the job only — VideoCount stays 0 until status returns video.url.
require.Equal(t, 0, result.ImageCount)
require.Empty(t, result.ImageSize)
require.Equal(t, 1, result.VideoCount)
require.Equal(t, 0, result.VideoCount)
require.Equal(t, VideoBillingResolution720P, result.VideoResolution)
require.Equal(t, 10, result.VideoDurationSeconds)
}
@@ -1383,7 +1384,7 @@ func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) {
require.Equal(t, http.MethodPost, upstream.lastReq.Method)
require.JSONEq(t, `{"model":"vendor-video-mutation","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`, string(upstream.lastBody))
require.Equal(t, "video-mutation-123", result.ResponseID)
require.Equal(t, 1, result.VideoCount)
require.Equal(t, 0, result.VideoCount)
require.Equal(t, 6, result.VideoDurationSeconds)
require.Equal(t, "grok-imagine-video", result.BillingModel)
require.Equal(t, "vendor-video-mutation", result.UpstreamModel)
@@ -664,6 +664,16 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i
return nil
}
func (c *stubGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c *stubGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c *stubGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) {
now := time.Now()
resetAt := now.Add(10 * time.Minute)
@@ -193,6 +193,16 @@ func (c *openAIWSStateStoreTimeoutProbeCache) DeleteSessionAccountID(ctx context
return nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c *openAIWSStateStoreTimeoutProbeCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
func TestOpenAIWSStateStore_RedisOpsUseShortTimeout(t *testing.T) {
probe := &openAIWSStateStoreTimeoutProbeCache{}
store := NewOpenAIWSStateStore(probe)
+11
View File
@@ -104,6 +104,17 @@ func (c StubGatewayCache) DeleteSessionAccountID(_ context.Context, _ int64, _ s
return nil
}
func (c StubGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
return nil
}
func (c StubGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
return nil, nil
}
func (c StubGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
return true, nil
}
// ============================================================
// StubSessionLimitCache — service.SessionLimitCache 的空实现
// ============================================================