From 87e6d9cf252da91402a8a0a1afa5df7763f5898f Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 24 Jun 2026 16:10:53 +0800 Subject: [PATCH] feat(videos): add model binding and propagation for video auth management - Introduced `model` field in `videoAuthBinding` for model tracking. - Enhanced binding store with `setWithModel` and `getBinding` methods for binding models. - Updated handlers to propagate execution models during video processing. - Added tests to verify model binding, retrieval, and propagation logic. --- .../handlers/openai/openai_videos_handlers.go | 73 +++++++++++++---- .../openai/openai_videos_handlers_test.go | 82 +++++++++++++++++++ 2 files changed, 137 insertions(+), 18 deletions(-) diff --git a/sdk/api/handlers/openai/openai_videos_handlers.go b/sdk/api/handlers/openai/openai_videos_handlers.go index 01b5ce6b..e891dbe2 100644 --- a/sdk/api/handlers/openai/openai_videos_handlers.go +++ b/sdk/api/handlers/openai/openai_videos_handlers.go @@ -55,6 +55,7 @@ type xaiVideoCreateMetadata struct { type videoAuthBinding struct { authID string + model string expiresAt time.Time } @@ -70,6 +71,10 @@ func newVideoAuthBindingStore() *videoAuthBindingStore { } func (s *videoAuthBindingStore) set(videoID string, authID string, ttl time.Duration) { + s.setWithModel(videoID, authID, "", ttl) +} + +func (s *videoAuthBindingStore) setWithModel(videoID string, authID string, model string, ttl time.Duration) { if s == nil { return } @@ -86,25 +91,34 @@ func (s *videoAuthBindingStore) set(videoID string, authID string, ttl time.Dura s.cleanupExpiredLocked(now) s.entries[videoID] = videoAuthBinding{ authID: authID, + model: strings.TrimSpace(model), expiresAt: now.Add(ttl), } s.mu.Unlock() } func (s *videoAuthBindingStore) get(videoID string) (string, bool) { - if s == nil { + binding, ok := s.getBinding(videoID) + if !ok { return "", false } + return binding.authID, true +} + +func (s *videoAuthBindingStore) getBinding(videoID string) (videoAuthBinding, bool) { + if s == nil { + return videoAuthBinding{}, false + } videoID = strings.TrimSpace(videoID) if videoID == "" { - return "", false + return videoAuthBinding{}, false } now := time.Now() s.mu.RLock() entry, ok := s.entries[videoID] s.mu.RUnlock() if !ok { - return "", false + return videoAuthBinding{}, false } if now.After(entry.expiresAt) { s.mu.Lock() @@ -112,9 +126,9 @@ func (s *videoAuthBindingStore) get(videoID string) (string, bool) { delete(s.entries, videoID) } s.mu.Unlock() - return "", false + return videoAuthBinding{}, false } - return entry.authID, true + return entry, true } func (s *videoAuthBindingStore) cleanupExpiredLocked(now time.Time) { @@ -276,11 +290,19 @@ func videoIDFromPayload(payload []byte) string { } func (h *OpenAIAPIHandler) bindVideoAuthIDFromPayload(payload []byte, authID string) { + h.bindVideoAuthIDAndModelFromPayload(payload, authID, strings.TrimSpace(gjson.GetBytes(payload, "model").String())) +} + +func (h *OpenAIAPIHandler) bindVideoAuthIDAndModelFromPayload(payload []byte, authID string, model string) { videoID := videoIDFromPayload(payload) if videoID == "" { return } - videoAuthBindings.set(videoID, authID, h.videoAuthBindingTTL()) + videoAuthBindings.setWithModel(videoID, authID, canonicalXAIVideosModel(model), h.videoAuthBindingTTL()) +} + +func (h *OpenAIAPIHandler) bindVideoAuthID(videoID string, authID string, model string) { + videoAuthBindings.setWithModel(videoID, authID, canonicalXAIVideosModel(model), h.videoAuthBindingTTL()) } func (h *OpenAIAPIHandler) contextWithVideoAuthBinding(ctx context.Context, videoID string) context.Context { @@ -290,6 +312,15 @@ func (h *OpenAIAPIHandler) contextWithVideoAuthBinding(ctx context.Context, vide return ctx } +func (h *OpenAIAPIHandler) modelWithVideoAuthBinding(videoID string, fallbackModel string) string { + if binding, ok := videoAuthBindings.getBinding(videoID); ok { + if model := strings.TrimSpace(binding.model); model != "" { + return model + } + } + return fallbackModel +} + func buildXAIVideosCreateRequest(rawJSON []byte, model string) ([]byte, xaiVideoCreateMetadata, error) { prompt := strings.TrimSpace(gjson.GetBytes(rawJSON, "prompt").String()) if prompt == "" { @@ -743,11 +774,12 @@ func (h *OpenAIAPIHandler) VideosRetrieve(c *gin.Context) { cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) selectedAuthID := "" cliCtx = h.contextWithVideoAuthBinding(cliCtx, videoID) + executionModel := h.modelWithVideoAuthBinding(videoID, defaultXAIVideosModel) cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { selectedAuthID = authID }) stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx) - resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, defaultXAIVideosModel, payload, "") + resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, payload, "") stopKeepAlive() if errMsg != nil { h.WriteErrorResponse(c, errMsg) @@ -767,7 +799,7 @@ func (h *OpenAIAPIHandler) VideosRetrieve(c *gin.Context) { return } - videoAuthBindings.set(videoID, selectedAuthID, h.videoAuthBindingTTL()) + h.bindVideoAuthID(videoID, selectedAuthID, executionModel) handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = c.Writer.Write(out) cliCancel(nil) @@ -805,11 +837,12 @@ func (h *OpenAIAPIHandler) VideosContent(c *gin.Context) { cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) selectedAuthID := "" cliCtx = h.contextWithVideoAuthBinding(cliCtx, videoID) + executionModel := h.modelWithVideoAuthBinding(videoID, defaultXAIVideosModel) cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { selectedAuthID = authID }) stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx) - resp, _, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, defaultXAIVideosModel, payload, "") + resp, _, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, payload, "") stopKeepAlive() if errMsg != nil { h.WriteErrorResponse(c, errMsg) @@ -821,7 +854,7 @@ func (h *OpenAIAPIHandler) VideosContent(c *gin.Context) { return } - videoAuthBindings.set(videoID, selectedAuthID, h.videoAuthBindingTTL()) + h.bindVideoAuthID(videoID, selectedAuthID, executionModel) contentURL, err := xaiVideoContentURLFromPayload(resp) if err != nil { errMsg := &interfaces.ErrorMessage{StatusCode: http.StatusBadGateway, Error: err} @@ -922,15 +955,17 @@ func (h *OpenAIAPIHandler) collectXAIVideosNative(c *gin.Context, rawJSON []byte cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background()) selectedAuthID := "" - if bindCreatedVideoAuth { - cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { - selectedAuthID = authID - }) - } else { - cliCtx = h.contextWithVideoAuthBinding(cliCtx, videoIDFromPayload(rawJSON)) + videoID := videoIDFromPayload(rawJSON) + executionModel := model + if !bindCreatedVideoAuth { + cliCtx = h.contextWithVideoAuthBinding(cliCtx, videoID) + executionModel = h.modelWithVideoAuthBinding(videoID, model) } + cliCtx = handlers.WithSelectedAuthIDCallback(cliCtx, func(authID string) { + selectedAuthID = authID + }) stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx) - resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, model, rawJSON, "") + resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, xaiVideosHandlerType, executionModel, rawJSON, "") stopKeepAlive() if errMsg != nil { h.WriteErrorResponse(c, errMsg) @@ -943,7 +978,9 @@ func (h *OpenAIAPIHandler) collectXAIVideosNative(c *gin.Context, rawJSON []byte } if bindCreatedVideoAuth { - h.bindVideoAuthIDFromPayload(resp, selectedAuthID) + h.bindVideoAuthIDAndModelFromPayload(resp, selectedAuthID, executionModel) + } else { + h.bindVideoAuthID(videoID, selectedAuthID, executionModel) } handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders) _, _ = c.Writer.Write(resp) diff --git a/sdk/api/handlers/openai/openai_videos_handlers_test.go b/sdk/api/handlers/openai/openai_videos_handlers_test.go index 8707fd96..52f6ca09 100644 --- a/sdk/api/handlers/openai/openai_videos_handlers_test.go +++ b/sdk/api/handlers/openai/openai_videos_handlers_test.go @@ -67,6 +67,7 @@ type videoAuthCaptureExecutor struct { requestID string contentURL string authIDs []string + models []string } func (e *videoAuthCaptureExecutor) Identifier() string { return "xai" } @@ -78,6 +79,7 @@ func (e *videoAuthCaptureExecutor) Execute(_ context.Context, auth *coreauth.Aut } e.mu.Lock() e.authIDs = append(e.authIDs, authID) + e.models = append(e.models, req.Model) e.mu.Unlock() requestID := strings.TrimSpace(gjson.GetBytes(req.Payload, "request_id").String()) @@ -116,6 +118,14 @@ func (e *videoAuthCaptureExecutor) AuthIDs() []string { return out } +func (e *videoAuthCaptureExecutor) Models() []string { + e.mu.Lock() + defer e.mu.Unlock() + out := make([]string, len(e.models)) + copy(out, e.models) + return out +} + func resetVideoAuthBindingsForTest(t *testing.T) { t.Helper() previous := videoAuthBindings @@ -142,6 +152,7 @@ func newVideoAuthBindingTestHandler(t *testing.T, executor *videoAuthCaptureExec t.Fatalf("manager.Register(%s): %v", authID, errRegister) } registry.GetGlobalRegistry().RegisterClient(authID, auth.Provider, []*registry.ModelInfo{{ID: defaultXAIVideosModel}}) + manager.RefreshSchedulerEntry(authID) } t.Cleanup(func() { for _, authID := range authIDs { @@ -723,6 +734,77 @@ func TestXAIVideosNativeCreateBindsRetrieveToSelectedAuth(t *testing.T) { } } +func TestXAIVideosNativeRetrieveUsesBoundModel(t *testing.T) { + resetVideoAuthBindingsForTest(t) + executor := &videoAuthCaptureExecutor{requestID: "video-xai-preview-bound"} + manager := coreauth.NewManager(nil, &coreauth.RoundRobinSelector{}, nil) + manager.RegisterExecutor(executor) + + authModels := []struct { + authID string + model string + }{ + {authID: "video-xai-preview-default-auth", model: defaultXAIVideosModel}, + {authID: "video-xai-preview-auth", model: xaiVideos15PreviewModel}, + } + for _, entry := range authModels { + auth := &coreauth.Auth{ + ID: entry.authID, + Provider: "xai", + Status: coreauth.StatusActive, + } + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("manager.Register(%s): %v", entry.authID, errRegister) + } + registry.GetGlobalRegistry().RegisterClient(entry.authID, auth.Provider, []*registry.ModelInfo{{ID: entry.model}}) + manager.RefreshSchedulerEntry(entry.authID) + } + t.Cleanup(func() { + for _, entry := range authModels { + registry.GetGlobalRegistry().UnregisterClient(entry.authID) + } + }) + + base := apihandlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager) + handler := NewOpenAIAPIHandler(base) + + createResp := performVideosEndpointRequest(t, http.MethodPost, xaiVideosGenerationsAPI, "application/json", strings.NewReader(`{"model":"grok-imagine-video-1.5-preview","prompt":"make a video"}`), handler.XAIVideosGenerations) + if createResp.Code != http.StatusOK { + t.Fatalf("create status = %d, want %d: %s", createResp.Code, http.StatusOK, createResp.Body.String()) + } + videoID := gjson.GetBytes(createResp.Body.Bytes(), "request_id").String() + if videoID != executor.requestID { + t.Fatalf("created request_id = %q, want %q", videoID, executor.requestID) + } + + retrieveResp := performVideosRouteRequest(t, http.MethodGet, videosPath+"/:request_id", videosPath+"/"+videoID, "", nil, handler.XAIVideosRetrieve) + if retrieveResp.Code != http.StatusOK { + t.Fatalf("retrieve status = %d, want %d: %s", retrieveResp.Code, http.StatusOK, retrieveResp.Body.String()) + } + + authIDs := executor.AuthIDs() + if len(authIDs) != 2 { + t.Fatalf("authIDs = %v, want two calls", authIDs) + } + if authIDs[0] != "video-xai-preview-auth" || authIDs[1] != authIDs[0] { + t.Fatalf("authIDs = %v, want both calls to use video-xai-preview-auth", authIDs) + } + models := executor.Models() + if len(models) != 2 { + t.Fatalf("models = %v, want two calls", models) + } + if models[0] != xaiVideos15PreviewModel || models[1] != xaiVideos15PreviewModel { + t.Fatalf("models = %v, want both calls to use %s", models, xaiVideos15PreviewModel) + } + binding, ok := videoAuthBindings.getBinding(videoID) + if !ok { + t.Fatal("video auth binding was not stored") + } + if binding.authID != "video-xai-preview-auth" || binding.model != xaiVideos15PreviewModel { + t.Fatalf("binding = {authID:%q model:%q}, want {authID:%q model:%q}", binding.authID, binding.model, "video-xai-preview-auth", xaiVideos15PreviewModel) + } +} + func TestVideoAuthBindingTTLUsesConfig(t *testing.T) { base := apihandlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{VideoResultAuthCacheTTL: "45m"}, nil) handler := NewOpenAIAPIHandler(base)