fix(kling): persist proxy fallback asset binding

Co-Authored-By: Codex <noreply@anthropic.com>

controller/kling_aiping_native.go

@@ -173,6 +173,10 @@ func KlingAipingNativeProxy(c *gin.Context) {
 			c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()})
 			return
 		}
+		if err := persistKlingAipingProxyBindingIfNeeded(c.GetInt("id"), group, route.BillingModel, channel); err != nil {
+			c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()})
+			return
+		}
 	}
 	if setupErr := middleware.SetupContextForSelectedChannel(c, channel, route.BillingModel); setupErr != nil {
 		c.JSON(setupErr.StatusCode, gin.H{"code": setupErr.GetErrorCode(), "message": setupErr.Error()})
@@ -207,6 +211,18 @@ func resolveKlingAipingBoundChannelForProxy(c *gin.Context, group, billingModel
 	return channel
 }
 
+func persistKlingAipingProxyBindingIfNeeded(userId int, group, billingModel string, channel *model.Channel) error {
+	group = strings.TrimSpace(group)
+	billingModel = strings.TrimSpace(billingModel)
+	if group == "" || group == "auto" || billingModel == "" || channel == nil {
+		return nil
+	}
+	if !service.IsUsableVideoAssetChannelForFamily(channel, group, billingModel, service.VideoAssetFamilyKling) {
+		return nil
+	}
+	return service.BindVideoAssetChannel(userId, group, channel, service.VideoAssetFamilyKling)
+}
+
 func readJSONPayload(c *gin.Context) (map[string]any, error) {
 	body, err := common.GetBodyStorage(c)
 	if err != nil {

controller/kling_aiping_native_test.go

@@ -4,6 +4,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"strings"
 	"testing"
 
@@ -16,6 +17,8 @@ import (
 	"github.com/QuantumNous/new-api/service"
 	"github.com/gin-gonic/gin"
 	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
 )
 
 func TestKlingAipingSubmitPreparationUsesModelNameBeforeModel(t *testing.T) {
@@ -172,6 +175,34 @@ func TestDoKlingAipingProxyRequestUsesSelectedContextKey(t *testing.T) {
 	require.Equal(t, "Bearer selected-key", gotAuth)
 }
 
+func TestKlingAipingNativeProxyPersistsFallbackBinding(t *testing.T) {
+	service.InitHttpClient()
+	db := setupKlingAipingNativeProxyDB(t)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"code":0,"message":"success","data":[]}`))
+	}))
+	defer upstream.Close()
+	createKlingAipingProxyChannelForTest(t, db, 59, "default", "proxy-key", upstream.URL)
+	createKlingAipingProxyAbilityForTest(t, db, "default", klingaiping.ModelKlingAdvancedElements, 59, true)
+
+	w := httptest.NewRecorder()
+	_, engine := gin.CreateTestContext(w)
+	engine.GET("/v1/general/advanced-presets-elements", func(c *gin.Context) {
+		c.Set("id", 10)
+		common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
+		KlingAipingNativeProxy(c)
+	})
+
+	engine.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/general/advanced-presets-elements", nil))
+
+	require.Equal(t, http.StatusOK, w.Code)
+	binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default")
+	require.NoError(t, err)
+	require.NotNil(t, binding)
+	require.Equal(t, 59, binding.ChannelId)
+}
+
 func TestCopyProxyResponseNormalizesMsgError(t *testing.T) {
 	w := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(w)
@@ -248,3 +279,85 @@ func TestParseKlingAipingPageBoundaries(t *testing.T) {
 	_, _, err = parseKlingAipingPage(c)
 	require.ErrorContains(t, err, "pageSize")
 }
+
+func setupKlingAipingNativeProxyDB(t *testing.T) *gorm.DB {
+	t.Helper()
+
+	oldDB := model.DB
+	oldSQLitePath := common.SQLitePath
+	oldMemoryCacheEnabled := common.MemoryCacheEnabled
+	oldIsMasterNode := common.IsMasterNode
+	oldUsingSQLite := common.UsingSQLite
+	oldUsingMySQL := common.UsingMySQL
+	oldUsingPostgreSQL := common.UsingPostgreSQL
+	oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN")
+
+	common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared"
+	common.MemoryCacheEnabled = false
+	common.IsMasterNode = false
+	common.UsingSQLite = false
+	common.UsingMySQL = false
+	common.UsingPostgreSQL = false
+	require.NoError(t, os.Setenv("SQL_DSN", "local"))
+	require.NoError(t, model.InitDB())
+
+	model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)})
+	db := model.DB
+	sqlDB, err := db.DB()
+	require.NoError(t, err)
+	sqlDB.SetMaxOpenConns(1)
+	require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.UserAssetChannel{}))
+
+	t.Cleanup(func() {
+		_ = sqlDB.Close()
+		model.DB = oldDB
+		common.SQLitePath = oldSQLitePath
+		common.MemoryCacheEnabled = oldMemoryCacheEnabled
+		common.IsMasterNode = oldIsMasterNode
+		common.UsingSQLite = oldUsingSQLite
+		common.UsingMySQL = oldUsingMySQL
+		common.UsingPostgreSQL = oldUsingPostgreSQL
+		if hadSQLDSN {
+			_ = os.Setenv("SQL_DSN", oldSQLDSN)
+		} else {
+			_ = os.Unsetenv("SQL_DSN")
+		}
+	})
+
+	return db
+}
+
+func createKlingAipingProxyChannelForTest(t *testing.T, db *gorm.DB, id int, group string, key string, baseURL string) {
+	t.Helper()
+
+	priority := int64(id)
+	weight := uint(10)
+	autoBan := 1
+	require.NoError(t, db.Create(&model.Channel{
+		Id:       id,
+		Type:     constant.ChannelTypeKlingAiping,
+		Key:      key,
+		Status:   common.ChannelStatusEnabled,
+		Name:     "kling-aiping-proxy",
+		Group:    group,
+		Models:   klingaiping.ModelKlingAdvancedElements,
+		BaseURL:  common.GetPointer(baseURL),
+		Priority: &priority,
+		Weight:   &weight,
+		AutoBan:  &autoBan,
+	}).Error)
+}
+
+func createKlingAipingProxyAbilityForTest(t *testing.T, db *gorm.DB, group string, modelName string, channelId int, enabled bool) {
+	t.Helper()
+
+	priority := int64(channelId)
+	require.NoError(t, db.Create(&model.Ability{
+		Group:     group,
+		Model:     modelName,
+		ChannelId: channelId,
+		Enabled:   enabled,
+		Priority:  &priority,
+		Weight:    10,
+	}).Error)
+}