mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-19 16:17:56 +00:00
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:
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 的空实现
|
||||
// ============================================================
|
||||
|
||||
Reference in New Issue
Block a user