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.
2026-11-08 17:36:47 +08:00 · 2026-06-24 16:10:53 +08:00
parent c20ac28a54
commit 87e6d9cf25
2 changed files with 137 additions and 18 deletions
@@ -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)
@@ -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)