diff --git a/.env.example b/.env.example index 02dfe80b5..0c5fbc247 100644 --- a/.env.example +++ b/.env.example @@ -235,6 +235,14 @@ PLUGIN_RUNTIME_MAX_BUFFER_SIZE=5242880 # plugin install/upgrade timeout in minutes PLUGIN_INSTALL_TIMEOUT=15 +# cache the tenant's model installation list in redis. Turn it on only where dify runs with +# PLUGIN_MODEL_PROVIDERS_CACHE_ENABLED=false, otherwise both sides hold the same payload +# and the redis footprint doubles. +PLUGIN_MODEL_INSTALLATIONS_CACHE_ENABLED=false + +# how long a tenant's cached model installation list survives without an explicit invalidation, in minutes +PLUGIN_MODEL_INSTALLATIONS_CACHE_TTL=1440 + # dify backwards invocation write timeout in milliseconds DIFY_BACKWARDS_INVOCATION_WRITE_TIMEOUT=5000 # dify backwards invocation read timeout in milliseconds diff --git a/internal/core/plugin_manager/manager.go b/internal/core/plugin_manager/manager.go index c3c76eac8..c14f91b95 100644 --- a/internal/core/plugin_manager/manager.go +++ b/internal/core/plugin_manager/manager.go @@ -2,6 +2,7 @@ package plugin_manager import ( "fmt" + "time" lru "github.com/hashicorp/golang-lru/v2" "github.com/langgenius/dify-cloud-kit/oss" @@ -16,6 +17,7 @@ import ( "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/plugin_packager/decoder" "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache/helper" "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" ) @@ -138,6 +140,10 @@ func (p *PluginManager) Launch(configuration *app.Config) { } cache.SetKeyPrefix(configuration.RedisKeyPrefix) + helper.SetModelInstallationsCacheEnabled(configuration.PluginModelInstallationsCacheEnabled) + helper.SetModelInstallationsCacheTTL( + time.Duration(configuration.PluginModelInstallationsCacheTTL) * time.Minute, + ) // init redis client if configuration.RedisUseSentinel { diff --git a/internal/service/manage_plugin.go b/internal/service/manage_plugin.go index 230c0b4a8..27c028c59 100644 --- a/internal/service/manage_plugin.go +++ b/internal/service/manage_plugin.go @@ -393,47 +393,11 @@ func ListTools(tenant_id string, page int, page_size int) *entities.Response { } func ListModels(tenant_id string, page int, page_size int) *entities.Response { - type AIModel struct { - models.AIModelInstallation // pointer to avoid deep copy - - Declaration *plugin_entities.ModelProviderDeclaration `json:"declaration"` - } - - providers, err := db.GetAll[models.AIModelInstallation]( - db.Equal("tenant_id", tenant_id), - db.Page(page, page_size), - ) - + data, err := helper.CombinedListModelInstallations(tenant_id, page, page_size) if err != nil { return exception.InternalServerError(err).ToResponse() } - data := make([]AIModel, 0, len(providers)) - - for _, provider := range providers { - uniqueIdentifier := plugin_entities.PluginUniqueIdentifier(provider.PluginUniqueIdentifier) - var runtimeType plugin_entities.PluginRuntimeType - if uniqueIdentifier.RemoteLike() { - runtimeType = plugin_entities.PLUGIN_RUNTIME_TYPE_REMOTE - } else { - runtimeType = plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL - } - - declaration, err := helper.CombinedGetPluginDeclaration( - uniqueIdentifier, - runtimeType, - ) - - if err != nil { - return exception.InternalServerError(err).ToResponse() - } - - data = append(data, AIModel{ - AIModelInstallation: provider, - Declaration: declaration.Model, - }) - } - return entities.NewSuccessResponse(data) } diff --git a/internal/types/app/config.go b/internal/types/app/config.go index 04c33de8f..1d71ada56 100644 --- a/internal/types/app/config.go +++ b/internal/types/app/config.go @@ -115,6 +115,11 @@ type Config struct { PluginInstalledPath string `envconfig:"PLUGIN_INSTALLED_PATH" validate:"required"` // where the plugin finally installed PluginPackageCachePath string `envconfig:"PLUGIN_PACKAGE_CACHE_PATH"` // where plugin packages stored + // cache the tenant's model installation list in redis, off unless dify's own model providers cache is disabled + PluginModelInstallationsCacheEnabled bool `envconfig:"PLUGIN_MODEL_INSTALLATIONS_CACHE_ENABLED"` + // how long a tenant's cached model installation list survives without an explicit invalidation, in minutes + PluginModelInstallationsCacheTTL int `envconfig:"PLUGIN_MODEL_INSTALLATIONS_CACHE_TTL"` + // request timeout PluginMaxExecutionTimeout int `envconfig:"PLUGIN_MAX_EXECUTION_TIMEOUT" validate:"required"` diff --git a/internal/types/app/default.go b/internal/types/app/default.go index 34f8b1241..749dd062d 100644 --- a/internal/types/app/default.go +++ b/internal/types/app/default.go @@ -40,6 +40,7 @@ func (config *Config) SetDefault() { setDefaultInt(&config.PluginLocalLaunchingConcurrent, 2) setDefaultInt(&config.PersistenceStorageMaxSize, 100*1024*1024) setDefaultString(&config.PluginPackageCachePath, "plugin_packages") + setDefaultInt(&config.PluginModelInstallationsCacheTTL, 1440) setDefaultString(&config.PythonInterpreterPath, "/usr/bin/python3") setDefaultInt(&config.PythonEnvInitTimeout, 120) setDefaultInt(&config.DifyInvocationWriteTimeout, 5000) diff --git a/internal/types/models/curd/atomic.go b/internal/types/models/curd/atomic.go index 3cae1624c..0788c74d4 100644 --- a/internal/types/models/curd/atomic.go +++ b/internal/types/models/curd/atomic.go @@ -29,6 +29,12 @@ func InstallPlugin( var pluginToBeReturns *models.Plugin var installationToBeReturns *models.PluginInstallation + // the cache is dropped even when the transaction fails, so a retry after a failed + // invalidation still clears it + if declaration.Model != nil { + defer helper.DeleteModelInstallationsCache(tenantId) + } + err := db.WithTransaction(func(tx *gorm.DB) error { pluginID := pluginUniqueIdentifier.PluginID() @@ -236,6 +242,10 @@ func UninstallPlugin( var pluginToBeReturns *models.Plugin var installationToBeReturns *models.PluginInstallation + if declaration.Model != nil { + defer helper.DeleteModelInstallationsCache(tenantId) + } + _, err := db.GetOne[models.PluginInstallation]( db.Equal("id", installationId), db.Equal("plugin_unique_identifier", pluginUniqueIdentifier.String()), @@ -428,6 +438,10 @@ func UpgradePlugin( ) (*UpgradePluginResponse, error) { var response UpgradePluginResponse + if originalDeclaration.Model != nil || newDeclaration.Model != nil { + defer helper.DeleteModelInstallationsCache(tenantId) + } + err := db.WithTransaction(func(tx *gorm.DB) error { installation, err := db.GetOne[models.PluginInstallation]( db.WithTransactionContext(tx), @@ -683,5 +697,6 @@ func UpgradePlugin( if _, err = cache.AutoDelete[models.PluginInstallation](pluginInstallationCacheKey); err != nil { return nil, err } + return &response, nil } diff --git a/internal/types/models/curd/model_installations_cache_test.go b/internal/types/models/curd/model_installations_cache_test.go new file mode 100644 index 000000000..3708388ed --- /dev/null +++ b/internal/types/models/curd/model_installations_cache_test.go @@ -0,0 +1,90 @@ +package curd + +import ( + "testing" + + "github.com/langgenius/dify-plugin-daemon/internal/db" + "github.com/langgenius/dify-plugin-daemon/internal/types/models" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache/helper" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + CACHE_TEST_TENANT_ID = "3d9e7c11-5b2a-4f8d-9e01-7a6b5c4d3e2f" + CACHE_TEST_PLUGIN_UNIQUE_IDENTIFIER = "acme/vertex-connector:0.4.1@0123456789abcdef" + CACHE_TEST_PLUGIN_ID = "acme/vertex-connector" + CACHE_TEST_PAGE_FIELD = "1:256" + CACHE_TEST_SENTINEL = `[{"provider":"stale-entry"}]` +) + +func cleanupInstallPluginCacheTest(t *testing.T) { + t.Helper() + t.Cleanup(func() { + _ = db.DeleteByCondition(&models.PluginInstallation{TenantID: CACHE_TEST_TENANT_ID}) + _ = db.DeleteByCondition(&models.AIModelInstallation{TenantID: CACHE_TEST_TENANT_ID}) + _ = db.DeleteByCondition(&models.Plugin{ + PluginUniqueIdentifier: CACHE_TEST_PLUGIN_UNIQUE_IDENTIFIER, + }) + _, _ = cache.Del(helper.ModelInstallationsCacheKey(CACHE_TEST_TENANT_ID)) + }) +} + +func seedModelInstallationsCache(t *testing.T) { + t.Helper() + require.NoError(t, cache.SetMapOneField( + helper.ModelInstallationsCacheKey(CACHE_TEST_TENANT_ID), + CACHE_TEST_PAGE_FIELD, + CACHE_TEST_SENTINEL, + )) + require.True(t, modelInstallationsCacheExists(t), "seed must land before the assertion is meaningful") +} + +func modelInstallationsCacheExists(t *testing.T) bool { + t.Helper() + exists, err := cache.Exist(helper.ModelInstallationsCacheKey(CACHE_TEST_TENANT_ID)) + require.NoError(t, err) + return exists == 1 +} + +func installCacheTestPlugin(declaration *plugin_entities.PluginDeclaration) error { + _, _, err := InstallPlugin( + CACHE_TEST_TENANT_ID, + plugin_entities.PluginUniqueIdentifier(CACHE_TEST_PLUGIN_UNIQUE_IDENTIFIER), + plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL, + declaration, + "test", + nil, + ) + return err +} + +// The invalidation must also run when the transaction fails, otherwise a retry after a +// failed Redis DEL can never clear the cache: it aborts on ErrPluginAlreadyInstalled first. +func TestInstallPluginInvalidatesModelInstallationsCacheOnRetry(t *testing.T) { + cleanupInstallPluginCacheTest(t) + + declaration := &plugin_entities.PluginDeclaration{ + Model: &plugin_entities.ModelProviderDeclaration{ + Provider: "declaration-provider", + }, + } + + seedModelInstallationsCache(t) + require.NoError(t, installCacheTestPlugin(declaration)) + assert.False(t, modelInstallationsCacheExists(t), "a successful install must drop the tenant cache") + + seedModelInstallationsCache(t) + require.ErrorIs(t, installCacheTestPlugin(declaration), ErrPluginAlreadyInstalled) + assert.False(t, modelInstallationsCacheExists(t), "a retry must still drop the tenant cache") +} + +func TestInstallPluginLeavesModelInstallationsCacheAloneWithoutModel(t *testing.T) { + cleanupInstallPluginCacheTest(t) + + seedModelInstallationsCache(t) + require.NoError(t, installCacheTestPlugin(&plugin_entities.PluginDeclaration{})) + assert.True(t, modelInstallationsCacheExists(t), "a plugin without a model must not invalidate") +} diff --git a/pkg/utils/cache/helper/keys.go b/pkg/utils/cache/helper/keys.go index 32fd92891..91057ca7c 100644 --- a/pkg/utils/cache/helper/keys.go +++ b/pkg/utils/cache/helper/keys.go @@ -14,6 +14,17 @@ func PluginInstallationCacheKey(pluginId, tenantId string) string { ) } +func ModelInstallationsCacheKey(tenantId string) string { + return strings.Join( + []string{ + "model_installations", + "tenant_id", + tenantId, + }, + ":", + ) +} + func EndpointCacheKey(hookId string) string { return strings.Join( []string{ diff --git a/pkg/utils/cache/helper/model_installations.go b/pkg/utils/cache/helper/model_installations.go new file mode 100644 index 000000000..42a8ea062 --- /dev/null +++ b/pkg/utils/cache/helper/model_installations.go @@ -0,0 +1,133 @@ +package helper + +import ( + "errors" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/langgenius/dify-plugin-daemon/internal/db" + "github.com/langgenius/dify-plugin-daemon/internal/types/models" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" +) + +const DEFAULT_MODEL_INSTALLATIONS_CACHE_TTL = 60 * time.Minute + +var ( + modelInstallationsCacheEnabled atomic.Bool + modelInstallationsCacheTTL atomic.Int64 +) + +func SetModelInstallationsCacheEnabled(enabled bool) { + modelInstallationsCacheEnabled.Store(enabled) +} + +func SetModelInstallationsCacheTTL(ttl time.Duration) { + if ttl <= 0 { + ttl = DEFAULT_MODEL_INSTALLATIONS_CACHE_TTL + } + modelInstallationsCacheTTL.Store(int64(ttl)) +} + +func currentModelInstallationsCacheTTL() time.Duration { + if ttl := modelInstallationsCacheTTL.Load(); ttl > 0 { + return time.Duration(ttl) + } + return DEFAULT_MODEL_INSTALLATIONS_CACHE_TTL +} + +type AIModelInstallationWithDeclaration struct { + models.AIModelInstallation + + Declaration *plugin_entities.ModelProviderDeclaration `json:"declaration"` +} + +func modelInstallationsCacheField(page int, pageSize int) string { + return strings.Join( + []string{ + strconv.Itoa(page), + strconv.Itoa(pageSize), + }, + ":", + ) +} + +func CombinedListModelInstallations( + tenantId string, + page int, + pageSize int, +) ([]AIModelInstallationWithDeclaration, error) { + if !modelInstallationsCacheEnabled.Load() { + return listModelInstallations(tenantId, page, pageSize) + } + + cacheKey := ModelInstallationsCacheKey(tenantId) + cacheField := modelInstallationsCacheField(page, pageSize) + + cached, err := cache.GetMapField[[]AIModelInstallationWithDeclaration](cacheKey, cacheField) + if err == nil { + return *cached, nil + } + if !errors.Is(err, cache.ErrNotFound) { + log.Warn("failed to read model installations cache", "key", cacheKey, "error", err) + } + + data, err := listModelInstallations(tenantId, page, pageSize) + if err != nil { + return nil, err + } + + if err := cache.SetMapOneField(cacheKey, cacheField, data); err != nil { + log.Warn("failed to store model installations cache", "key", cacheKey, "error", err) + } else if _, err := cache.Expire(cacheKey, currentModelInstallationsCacheTTL()); err != nil { + log.Warn("failed to set model installations cache expiry", "key", cacheKey, "error", err) + } + + return data, nil +} + +func listModelInstallations( + tenantId string, + page int, + pageSize int, +) ([]AIModelInstallationWithDeclaration, error) { + installations, err := db.GetAll[models.AIModelInstallation]( + db.Equal("tenant_id", tenantId), + db.Page(page, pageSize), + ) + if err != nil { + return nil, err + } + + data := make([]AIModelInstallationWithDeclaration, 0, len(installations)) + + for _, installation := range installations { + uniqueIdentifier := plugin_entities.PluginUniqueIdentifier(installation.PluginUniqueIdentifier) + runtimeType := plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL + if uniqueIdentifier.RemoteLike() { + runtimeType = plugin_entities.PLUGIN_RUNTIME_TYPE_REMOTE + } + + declaration, err := CombinedGetPluginDeclaration(uniqueIdentifier, runtimeType) + if err != nil { + return nil, err + } + + data = append(data, AIModelInstallationWithDeclaration{ + AIModelInstallation: installation, + Declaration: declaration.Model, + }) + } + + return data, nil +} + +func DeleteModelInstallationsCache(tenantId string) { + cacheKey := ModelInstallationsCacheKey(tenantId) + if _, err := cache.Del(cacheKey); err != nil { + log.Warn("failed to clear model installations cache", "key", cacheKey, "error", err) + } +} diff --git a/pkg/utils/cache/helper/model_installations_test.go b/pkg/utils/cache/helper/model_installations_test.go new file mode 100644 index 000000000..4eb23b5b0 --- /dev/null +++ b/pkg/utils/cache/helper/model_installations_test.go @@ -0,0 +1,149 @@ +package helper + +import ( + "strings" + "testing" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/langgenius/dify-plugin-daemon/internal/db" + "github.com/langgenius/dify-plugin-daemon/internal/types/models" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +const ( + TEST_TENANT_ID = "8f2c1b6e-3a4d-4f5e-9a7b-1c2d3e4f5a6b" + // deliberately all different, so confusing any two of them fails a test + TEST_PLUGIN_UNIQUE_IDENTIFIER = "acme/vertex-connector:0.4.1@0123456789abcdef" + TEST_PLUGIN_ID = "acme/vertex-connector" + TEST_ROW_PROVIDER = "row-provider" + TEST_DECLARATION_PROVIDER = "declaration-provider" +) + +func setupModelInstallationsTest(t *testing.T) { + t.Helper() + + SetModelInstallationsCacheEnabled(true) + t.Cleanup(func() { + SetModelInstallationsCacheEnabled(false) + }) + + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = cache.Close() + }) + + gormDB, err := gorm.Open( + sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), + &gorm.Config{}, + ) + require.NoError(t, err) + require.NoError(t, gormDB.AutoMigrate( + &models.AIModelInstallation{}, + &models.PluginDeclaration{}, + )) + db.DifyPluginDB = gormDB + t.Cleanup(func() { + if sqlDB, err := gormDB.DB(); err == nil { + _ = sqlDB.Close() + } + }) + + require.NoError(t, gormDB.Create(&models.PluginDeclaration{ + PluginUniqueIdentifier: TEST_PLUGIN_UNIQUE_IDENTIFIER, + PluginID: TEST_PLUGIN_ID, + Declaration: plugin_entities.PluginDeclaration{ + Model: &plugin_entities.ModelProviderDeclaration{ + Provider: TEST_DECLARATION_PROVIDER, + }, + }, + }).Error) + + require.NoError(t, gormDB.Create(&models.AIModelInstallation{ + PluginID: TEST_PLUGIN_ID, + PluginUniqueIdentifier: TEST_PLUGIN_UNIQUE_IDENTIFIER, + TenantID: TEST_TENANT_ID, + Provider: TEST_ROW_PROVIDER, + }).Error) + + t.Cleanup(func() { + DeletePluginDeclarationCache( + plugin_entities.PluginUniqueIdentifier(TEST_PLUGIN_UNIQUE_IDENTIFIER), + plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL, + ) + }) +} + +func deleteInstallationRows(t *testing.T) { + t.Helper() + require.NoError(t, db.DifyPluginDB. + Where("tenant_id = ?", TEST_TENANT_ID). + Delete(&models.AIModelInstallation{}).Error) +} + +func TestCombinedListModelInstallationsServesCachedPageAfterRowsChange(t *testing.T) { + setupModelInstallationsTest(t) + + first, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + require.Len(t, first, 1) + assert.Equal(t, TEST_ROW_PROVIDER, first[0].Provider) + assert.Equal(t, TEST_PLUGIN_ID, first[0].PluginID) + require.NotNil(t, first[0].Declaration) + assert.Equal(t, TEST_DECLARATION_PROVIDER, first[0].Declaration.Provider) + + deleteInstallationRows(t) + + cached, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + assert.Len(t, cached, 1, "a cache hit must not reach the database") + + DeleteModelInstallationsCache(TEST_TENANT_ID) + + fresh, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + assert.Empty(t, fresh) +} + +func TestCombinedListModelInstallationsBypassesCacheWhenDisabled(t *testing.T) { + setupModelInstallationsTest(t) + SetModelInstallationsCacheEnabled(false) + + first, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + require.Len(t, first, 1) + + exists, err := cache.Exist(ModelInstallationsCacheKey(TEST_TENANT_ID)) + require.NoError(t, err) + assert.Zero(t, exists, "a disabled cache must not write the tenant key") + + deleteInstallationRows(t) + + fresh, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + assert.Empty(t, fresh, "a disabled cache must always reach the database") +} + +func TestCombinedListModelInstallationsCachesEachPageSeparately(t *testing.T) { + setupModelInstallationsTest(t) + + firstPage, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + require.Len(t, firstPage, 1) + + secondPage, err := CombinedListModelInstallations(TEST_TENANT_ID, 2, 256) + require.NoError(t, err) + assert.Empty(t, secondPage, "page 2 must not be served the page 1 entry") + + DeleteModelInstallationsCache(TEST_TENANT_ID) + deleteInstallationRows(t) + + afterInvalidation, err := CombinedListModelInstallations(TEST_TENANT_ID, 1, 256) + require.NoError(t, err) + assert.Empty(t, afterInvalidation, "one delete must drop every cached page") +} diff --git a/pkg/utils/parser/json.go b/pkg/utils/parser/json.go index 7a98a68e0..aa2431443 100644 --- a/pkg/utils/parser/json.go +++ b/pkg/utils/parser/json.go @@ -18,11 +18,9 @@ func UnmarshalJsonBytes[T any](data []byte) (T, error) { return result, err } - // skip validate if T is a map - typ := reflect.TypeOf(result) - if typ.Kind() == reflect.Map { - return result, nil - } else if typ.Kind() == reflect.String { + // the entity validator only accepts structs, so kinds it cannot handle skip validation + switch reflect.TypeOf(result).Kind() { + case reflect.Map, reflect.String, reflect.Slice: return result, nil }