From bde79b7259fcc1479fcaea5176d665068a64fc07 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Fri, 10 Jul 2026 09:29:05 +0800 Subject: [PATCH] fix(kling): persist proxy fallback asset binding Co-Authored-By: Codex --- controller/kling_aiping_native.go | 16 ++++ controller/kling_aiping_native_test.go | 113 +++++++++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/controller/kling_aiping_native.go b/controller/kling_aiping_native.go index f38fe80..01e21d3 100644 --- a/controller/kling_aiping_native.go +++ b/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 { diff --git a/controller/kling_aiping_native_test.go b/controller/kling_aiping_native_test.go index 5322f71..c733275 100644 --- a/controller/kling_aiping_native_test.go +++ b/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) +}