From 8ae307d9bc6a5b7a54b67ba621e6dcc0cf7a56b3 Mon Sep 17 00:00:00 2001 From: GareArc Date: Mon, 27 Jul 2026 01:42:40 -0700 Subject: [PATCH 1/2] fix(codegen): restore the plugin daemon output path The generator still wrote to internal/core/plugin_daemon, which was renamed to io_tunnel, so `go run cmd/codegen/main.go` failed partway through and left the tree half-regenerated. Regenerating with the corrected path also settles import ordering in the files that were reformatted by hand after the last run. --- internal/core/io_tunnel/datasource.gen.go | 2 +- internal/core/io_tunnel/dynamic_parameter.gen.go | 2 +- internal/core/io_tunnel/oauth.gen.go | 2 +- internal/core/io_tunnel/tool.gen.go | 2 +- internal/core/io_tunnel/trigger.gen.go | 2 +- internal/server/controllers/generator/generator.go | 2 +- internal/service/datasource.gen.go | 2 +- internal/service/dynamic_parameter.gen.go | 2 +- internal/service/oauth.gen.go | 2 +- internal/service/tool.gen.go | 2 +- internal/service/trigger.gen.go | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/internal/core/io_tunnel/datasource.gen.go b/internal/core/io_tunnel/datasource.gen.go index a92fdf951..0cffe3d50 100644 --- a/internal/core/io_tunnel/datasource.gen.go +++ b/internal/core/io_tunnel/datasource.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/datasource_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func DatasourceValidateCredentials( diff --git a/internal/core/io_tunnel/dynamic_parameter.gen.go b/internal/core/io_tunnel/dynamic_parameter.gen.go index 28968a0dd..03b51dcdc 100644 --- a/internal/core/io_tunnel/dynamic_parameter.gen.go +++ b/internal/core/io_tunnel/dynamic_parameter.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/dynamic_select_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func FetchDynamicParameterOptions( diff --git a/internal/core/io_tunnel/oauth.gen.go b/internal/core/io_tunnel/oauth.gen.go index 9b1621998..b652fd29c 100644 --- a/internal/core/io_tunnel/oauth.gen.go +++ b/internal/core/io_tunnel/oauth.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/oauth_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func GetAuthorizationURL( diff --git a/internal/core/io_tunnel/tool.gen.go b/internal/core/io_tunnel/tool.gen.go index 080009814..965777985 100644 --- a/internal/core/io_tunnel/tool.gen.go +++ b/internal/core/io_tunnel/tool.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" "github.com/langgenius/dify-plugin-daemon/pkg/entities/tool_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeTool( diff --git a/internal/core/io_tunnel/trigger.gen.go b/internal/core/io_tunnel/trigger.gen.go index 4941bd7e9..48ec208a6 100644 --- a/internal/core/io_tunnel/trigger.gen.go +++ b/internal/core/io_tunnel/trigger.gen.go @@ -4,8 +4,8 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func TriggerInvokeEvent( diff --git a/internal/server/controllers/generator/generator.go b/internal/server/controllers/generator/generator.go index 54756a64e..0449dd744 100644 --- a/internal/server/controllers/generator/generator.go +++ b/internal/server/controllers/generator/generator.go @@ -114,7 +114,7 @@ func GeneratePluginDaemon(accessType access_types.PluginAccessType, dispatchers tmpl := template.Must(template.New("pluginDaemon").Parse(pluginDaemonTemplate)) // Create output file - outputPath := filepath.Join("internal", "core", "plugin_daemon", strings.ToLower(string(accessType))+".gen.go") + outputPath := filepath.Join("internal", "core", "io_tunnel", strings.ToLower(string(accessType))+".gen.go") // Execute template var buf strings.Builder diff --git a/internal/service/datasource.gen.go b/internal/service/datasource.gen.go index deee6ee1d..b9a5c9b7e 100644 --- a/internal/service/datasource.gen.go +++ b/internal/service/datasource.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/datasource_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func DatasourceValidateCredentials( diff --git a/internal/service/dynamic_parameter.gen.go b/internal/service/dynamic_parameter.gen.go index 7f28cb59b..35b8e0c9b 100644 --- a/internal/service/dynamic_parameter.gen.go +++ b/internal/service/dynamic_parameter.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/dynamic_select_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func FetchDynamicParameterOptions( diff --git a/internal/service/oauth.gen.go b/internal/service/oauth.gen.go index 19d426582..ddddcc09f 100644 --- a/internal/service/oauth.gen.go +++ b/internal/service/oauth.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/oauth_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func GetAuthorizationURL( diff --git a/internal/service/tool.gen.go b/internal/service/tool.gen.go index 33dd4f23c..200198cef 100644 --- a/internal/service/tool.gen.go +++ b/internal/service/tool.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" "github.com/langgenius/dify-plugin-daemon/pkg/entities/tool_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeTool( diff --git a/internal/service/trigger.gen.go b/internal/service/trigger.gen.go index 8b73133d3..d7798efa5 100644 --- a/internal/service/trigger.gen.go +++ b/internal/service/trigger.gen.go @@ -7,9 +7,9 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func TriggerInvokeEvent( From f6afd392e3ea408f2ce260246e4ccd44e82dfbdd Mon Sep 17 00:00:00 2001 From: GareArc Date: Mon, 27 Jul 2026 01:42:50 -0700 Subject: [PATCH 2/2] feat(cache): cache model schemas in the daemon Dify caches model schemas under a key built from the plugin id with no version, so an enterprise-side upgrade leaves stale schemas behind and Dify has no way to learn about it. Cache them here instead, keyed on the plugin unique identifier, which pins version and checksum. An upgrade lands on a different key, so a stale entry is unreachable and nothing needs to invalidate it. GetAIModelSchema moves out of the generated dispatchers, following the InvokeAgentStrategy precedent, so a cache hit can answer without creating a session or invoking the plugin. Remote and debugging plugins keep one identifier across code changes and are never cached. TTL defaults to 60 minutes via PLUGIN_MODEL_SCHEMA_CACHE_TTL; since the key is content-addressed it only bounds memory. --- .env.example | 3 + internal/core/io_tunnel/model.gen.go | 15 +-- internal/core/plugin_manager/manager.go | 5 + .../controllers/definitions/definitions.go | 22 ++--- internal/server/controllers/model.gen.go | 13 --- internal/server/controllers/model_schema.go | 22 +++++ internal/server/http_server.gen.go | 1 - internal/server/http_server.go | 3 + internal/service/model.gen.go | 23 +---- internal/service/model_schema.go | 88 +++++++++++++++++ internal/types/app/config.go | 3 + internal/types/app/default.go | 1 + pkg/utils/cache/helper/model_schema.go | 93 ++++++++++++++++++ pkg/utils/cache/helper/model_schema_test.go | 98 +++++++++++++++++++ 14 files changed, 331 insertions(+), 59 deletions(-) create mode 100644 internal/server/controllers/model_schema.go create mode 100644 internal/service/model_schema.go create mode 100644 pkg/utils/cache/helper/model_schema.go create mode 100644 pkg/utils/cache/helper/model_schema_test.go diff --git a/.env.example b/.env.example index 02dfe80b5..4263fd893 100644 --- a/.env.example +++ b/.env.example @@ -235,6 +235,9 @@ PLUGIN_RUNTIME_MAX_BUFFER_SIZE=5242880 # plugin install/upgrade timeout in minutes PLUGIN_INSTALL_TIMEOUT=15 +# how long a cached model schema is kept, in minutes +PLUGIN_MODEL_SCHEMA_CACHE_TTL=60 + # 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/io_tunnel/model.gen.go b/internal/core/io_tunnel/model.gen.go index ba4d2ea8a..3607ebe24 100644 --- a/internal/core/io_tunnel/model.gen.go +++ b/internal/core/io_tunnel/model.gen.go @@ -4,9 +4,9 @@ package io_tunnel import ( "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeLLM( @@ -177,16 +177,3 @@ func ValidateModelCredentials( 1, ) } - -func GetAIModelSchema( - session *session_manager.Session, - request *requests.RequestGetAIModelSchema, -) ( - *stream.Stream[model_entities.GetModelSchemasResponse], error, -) { - return GenericInvokePlugin[requests.RequestGetAIModelSchema, model_entities.GetModelSchemasResponse]( - session, - request, - 1, - ) -} diff --git a/internal/core/plugin_manager/manager.go b/internal/core/plugin_manager/manager.go index c3c76eac8..c27377a3f 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,9 @@ func (p *PluginManager) Launch(configuration *app.Config) { } cache.SetKeyPrefix(configuration.RedisKeyPrefix) + helper.SetModelSchemaCacheTTL( + time.Duration(configuration.PluginModelSchemaCacheTTL) * time.Minute, + ) // init redis client if configuration.RedisUseSentinel { diff --git a/internal/server/controllers/definitions/definitions.go b/internal/server/controllers/definitions/definitions.go index 9f3b3f134..8efe2281b 100644 --- a/internal/server/controllers/definitions/definitions.go +++ b/internal/server/controllers/definitions/definitions.go @@ -203,17 +203,17 @@ var PluginDispatchers = []PluginDispatcher{ BufferSize: 1, Path: "/model/validate_model_credentials", }, - { - Name: "GetAIModelSchema", - RequestType: requests.RequestGetAIModelSchema{}, - ResponseType: model_entities.GetModelSchemasResponse{}, - AccessType: access_types.PLUGIN_ACCESS_TYPE_MODEL, - AccessAction: access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, - AccessTypeString: "access_types.PLUGIN_ACCESS_TYPE_MODEL", - AccessActionString: "access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS", - BufferSize: 1, - Path: "/model/schema", - }, + // { // No need to implement this for now, it is cached in the model schema service + // Name: "GetAIModelSchema", + // RequestType: requests.RequestGetAIModelSchema{}, + // ResponseType: model_entities.GetModelSchemasResponse{}, + // AccessType: access_types.PLUGIN_ACCESS_TYPE_MODEL, + // AccessAction: access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, + // AccessTypeString: "access_types.PLUGIN_ACCESS_TYPE_MODEL", + // AccessActionString: "access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS", + // BufferSize: 1, + // Path: "/model/schema", + // }, // { // No need to implement this for now, it has its special implementation in the agent service // Name: "InvokeAgentStrategy", // RequestType: requests.RequestInvokeAgentStrategy{}, diff --git a/internal/server/controllers/model.gen.go b/internal/server/controllers/model.gen.go index 7fb488937..e7e5d3c30 100644 --- a/internal/server/controllers/model.gen.go +++ b/internal/server/controllers/model.gen.go @@ -178,16 +178,3 @@ func ValidateModelCredentials(config *app.Config) gin.HandlerFunc { ) } } - -func GetAIModelSchema(config *app.Config) gin.HandlerFunc { - type request = plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema] - - return func(c *gin.Context) { - BindPluginDispatchRequest( - c, - func(itr request) { - service.GetAIModelSchema(&itr, c, config.PluginMaxExecutionTimeout) - }, - ) - } -} diff --git a/internal/server/controllers/model_schema.go b/internal/server/controllers/model_schema.go new file mode 100644 index 000000000..3d828d043 --- /dev/null +++ b/internal/server/controllers/model_schema.go @@ -0,0 +1,22 @@ +package controllers + +import ( + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/internal/service" + "github.com/langgenius/dify-plugin-daemon/internal/types/app" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" +) + +func GetAIModelSchema(config *app.Config) gin.HandlerFunc { + type request = plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema] + + return func(c *gin.Context) { + BindPluginDispatchRequest( + c, + func(itr request) { + service.GetAIModelSchema(&itr, c, config.PluginMaxExecutionTimeout) + }, + ) + } +} diff --git a/internal/server/http_server.gen.go b/internal/server/http_server.gen.go index 4392b0121..9e7d4ed81 100644 --- a/internal/server/http_server.gen.go +++ b/internal/server/http_server.gen.go @@ -24,7 +24,6 @@ func (app *App) setupGeneratedRoutes(group *gin.RouterGroup, config *app.Config) group.POST("/moderation/invoke", controllers.InvokeModeration(config)) group.POST("/model/validate_provider_credentials", controllers.ValidateProviderCredentials(config)) group.POST("/model/validate_model_credentials", controllers.ValidateModelCredentials(config)) - group.POST("/model/schema", controllers.GetAIModelSchema(config)) group.POST("/oauth/get_authorization_url", controllers.GetAuthorizationURL(config)) group.POST("/oauth/get_credentials", controllers.GetCredentials(config)) group.POST("/oauth/refresh_credentials", controllers.RefreshCredentials(config)) diff --git a/internal/server/http_server.go b/internal/server/http_server.go index 925cbdce3..cae40de09 100644 --- a/internal/server/http_server.go +++ b/internal/server/http_server.go @@ -121,6 +121,7 @@ func (app *App) pluginDispatchGroup(group *gin.RouterGroup, config *app.Config) group.Use(app.InitClusterID()) group.POST("/agent_strategy/invoke", controllers.InvokeAgentStrategy(config)) + group.POST("/model/schema", controllers.GetAIModelSchema(config)) group.POST("/model/polling/start", controllers.StartPolling()) group.POST("/model/polling/check", controllers.CheckPolling()) @@ -245,6 +246,8 @@ func (app *App) invokeGroup(group *gin.RouterGroup, config *app.Config) { dispatchGroup.POST("/agent_strategy/invoke", controllers.InvokeAgentStrategy(config)) + dispatchGroup.POST("/model/schema", + controllers.GetAIModelSchema(config)) app.setupGeneratedRoutes(dispatchGroup, config) } diff --git a/internal/service/model.gen.go b/internal/service/model.gen.go index 219580f24..d77758619 100644 --- a/internal/service/model.gen.go +++ b/internal/service/model.gen.go @@ -7,10 +7,10 @@ import ( "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" - "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) func InvokeLLM( @@ -74,7 +74,7 @@ func InvokeMultimodalEmbedding( return io_tunnel.InvokeMultimodalEmbedding(session, &r.Data) }, access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_INVOKE_MULTIMODAL_EMBEDDING, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_TEXT_EMBEDDING, r, ctx, max_timeout_seconds, @@ -125,7 +125,7 @@ func InvokeMultimodalRerank( return io_tunnel.InvokeMultimodalRerank(session, &r.Data) }, access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_INVOKE_MULTIMODAL_RERANK, + access_types.PLUGIN_ACCESS_ACTION_INVOKE_RERANK, r, ctx, max_timeout_seconds, @@ -233,20 +233,3 @@ func ValidateModelCredentials( max_timeout_seconds, ) } - -func GetAIModelSchema( - r *plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema], - ctx *gin.Context, - max_timeout_seconds int, -) { - baseSSEWithSession( - func(session *session_manager.Session) (*stream.Stream[model_entities.GetModelSchemasResponse], error) { - return io_tunnel.GetAIModelSchema(session, &r.Data) - }, - access_types.PLUGIN_ACCESS_TYPE_MODEL, - access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, - r, - ctx, - max_timeout_seconds, - ) -} diff --git a/internal/service/model_schema.go b/internal/service/model_schema.go new file mode 100644 index 000000000..6efdeac91 --- /dev/null +++ b/internal/service/model_schema.go @@ -0,0 +1,88 @@ +package service + +import ( + "github.com/gin-gonic/gin" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel" + "github.com/langgenius/dify-plugin-daemon/internal/core/io_tunnel/access_types" + "github.com/langgenius/dify-plugin-daemon/internal/core/session_manager" + "github.com/langgenius/dify-plugin-daemon/pkg/entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/requests" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache/helper" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" +) + +func GetAIModelSchema( + r *plugin_entities.InvokePluginRequest[requests.RequestGetAIModelSchema], + ctx *gin.Context, + max_timeout_seconds int, +) { + identity := helper.ModelSchemaCacheIdentity{ + UniqueIdentifier: r.UniqueIdentifier, + ModelType: string(r.Data.ModelType), + Model: r.Data.Model, + CredentialType: r.Data.Credentials.CredentialType, + Credentials: r.Data.Credentials.Credentials, + } + + if cached := helper.GetCachedModelSchema(identity); cached != nil { + writeSingleSSEChunk(ctx, *cached) + return + } + + baseSSEWithSession( + func(session *session_manager.Session) (*stream.Stream[model_entities.GetModelSchemasResponse], error) { + response, err := io_tunnel.GenericInvokePlugin[ + requests.RequestGetAIModelSchema, + model_entities.GetModelSchemasResponse, + ](session, &r.Data, 1) + if err != nil { + return nil, err + } + + cacheModelSchemaStream(response, identity) + return response, nil + }, + access_types.PLUGIN_ACCESS_TYPE_MODEL, + access_types.PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS, + r, + ctx, + max_timeout_seconds, + ) +} + +// The plugin answers with a single chunk. Anything else, including an error mid-stream, +// leaves the cache untouched rather than storing a partial answer. +func cacheModelSchemaStream( + response *stream.Stream[model_entities.GetModelSchemasResponse], + identity helper.ModelSchemaCacheIdentity, +) { + var chunks []model_entities.GetModelSchemasResponse + + response.Filter(func(chunk model_entities.GetModelSchemasResponse) error { + chunks = append(chunks, chunk) + return nil + }) + + response.OnClose(func() { + if len(chunks) != 1 || chunks[0].ModelSchema == nil { + return + } + helper.StoreModelSchema(identity, chunks[0]) + }) +} + +// Mirrors the wire format baseSSEService writes, so a cache hit is indistinguishable to the +// caller from a plugin invocation. +func writeSingleSSEChunk(ctx *gin.Context, chunk model_entities.GetModelSchemasResponse) { + writer := ctx.Writer + writer.WriteHeader(200) + writer.Header().Set("Content-Type", "text/event-stream") + + writer.Write([]byte("data: ")) + writer.Write(parser.MarshalJsonBytes(entities.NewSuccessResponse(chunk))) + writer.Write([]byte("\n\n")) + writer.Flush() +} diff --git a/internal/types/app/config.go b/internal/types/app/config.go index 04c33de8f..5582af55e 100644 --- a/internal/types/app/config.go +++ b/internal/types/app/config.go @@ -115,6 +115,9 @@ 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 + // how long a cached model schema is kept, in minutes; the key pins the plugin version, so this only bounds memory + PluginModelSchemaCacheTTL int `envconfig:"PLUGIN_MODEL_SCHEMA_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..1def38e12 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.PluginModelSchemaCacheTTL, 60) setDefaultString(&config.PythonInterpreterPath, "/usr/bin/python3") setDefaultInt(&config.PythonEnvInitTimeout, 120) setDefaultInt(&config.DifyInvocationWriteTimeout, 5000) diff --git a/pkg/utils/cache/helper/model_schema.go b/pkg/utils/cache/helper/model_schema.go new file mode 100644 index 000000000..95e20acd4 --- /dev/null +++ b/pkg/utils/cache/helper/model_schema.go @@ -0,0 +1,93 @@ +package helper + +import ( + "crypto/sha256" + "encoding/hex" + "strings" + "sync/atomic" + "time" + + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "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" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" +) + +const DEFAULT_MODEL_SCHEMA_CACHE_TTL = 60 * time.Minute + +var modelSchemaCacheTTL atomic.Int64 + +func SetModelSchemaCacheTTL(ttl time.Duration) { + if ttl <= 0 { + ttl = DEFAULT_MODEL_SCHEMA_CACHE_TTL + } + modelSchemaCacheTTL.Store(int64(ttl)) +} + +func currentModelSchemaCacheTTL() time.Duration { + if ttl := modelSchemaCacheTTL.Load(); ttl > 0 { + return time.Duration(ttl) + } + return DEFAULT_MODEL_SCHEMA_CACHE_TTL +} + +type ModelSchemaCacheIdentity struct { + UniqueIdentifier plugin_entities.PluginUniqueIdentifier + ModelType string + Model string + CredentialType string + Credentials map[string]any +} + +// Remote and debugging plugins keep one identifier across code changes, so their schemas +// are never cached. +func (i ModelSchemaCacheIdentity) cacheable() bool { + return !i.UniqueIdentifier.RemoteLike() +} + +func (i ModelSchemaCacheIdentity) cacheKey() string { + digest := sha256.Sum256(parser.MarshalJsonBytes(map[string]any{ + "credential_type": i.CredentialType, + "credentials": i.Credentials, + })) + + return strings.Join( + []string{ + "model_schema", + i.UniqueIdentifier.String(), + i.ModelType, + i.Model, + hex.EncodeToString(digest[:]), + }, + ":", + ) +} + +// The key pins the plugin version and checksum, so an upgrade lands on a different key and +// a stale entry is unreachable. Nothing invalidates this cache. +func GetCachedModelSchema(identity ModelSchemaCacheIdentity) *model_entities.GetModelSchemasResponse { + if !identity.cacheable() { + return nil + } + + schema, err := cache.Get[model_entities.GetModelSchemasResponse](identity.cacheKey()) + if err != nil { + if err != cache.ErrNotFound { + log.Warn("failed to read model schema cache", "model", identity.Model, "error", err) + } + return nil + } + + return schema +} + +func StoreModelSchema(identity ModelSchemaCacheIdentity, schema model_entities.GetModelSchemasResponse) { + if !identity.cacheable() { + return + } + + if err := cache.Store(identity.cacheKey(), schema, currentModelSchemaCacheTTL()); err != nil { + log.Warn("failed to store model schema cache", "model", identity.Model, "error", err) + } +} diff --git a/pkg/utils/cache/helper/model_schema_test.go b/pkg/utils/cache/helper/model_schema_test.go new file mode 100644 index 000000000..4174f565a --- /dev/null +++ b/pkg/utils/cache/helper/model_schema_test.go @@ -0,0 +1,98 @@ +package helper + +import ( + "testing" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities" + "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" +) + +const ( + SCHEMA_TEST_IDENTIFIER_V1 = "acme/vertex-connector:0.4.1@0123456789abcdef" + SCHEMA_TEST_IDENTIFIER_V2 = "acme/vertex-connector:0.5.0@fedcba9876543210" + SCHEMA_TEST_REMOTE = "5f8d1c2e-9a3b-4c6d-8e7f-0a1b2c3d4e5f/debug-plugin:0.0.1@aaaabbbbccccdddd" + SCHEMA_TEST_MODEL_TYPE = "llm" + SCHEMA_TEST_MODEL = "gemini-2.5-pro" +) + +func setupModelSchemaTest(t *testing.T) { + t.Helper() + + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = cache.Close() + }) +} + +func schemaTestIdentity(identifier string, credentials map[string]any) ModelSchemaCacheIdentity { + return ModelSchemaCacheIdentity{ + UniqueIdentifier: plugin_entities.PluginUniqueIdentifier(identifier), + ModelType: SCHEMA_TEST_MODEL_TYPE, + Model: SCHEMA_TEST_MODEL, + CredentialType: "api-key", + Credentials: credentials, + } +} + +func schemaWithLabel(label string) model_entities.GetModelSchemasResponse { + return model_entities.GetModelSchemasResponse{ + ModelSchema: &plugin_entities.ModelDeclaration{ + Model: label, + }, + } +} + +func TestModelSchemaCacheRoundTrip(t *testing.T) { + setupModelSchemaTest(t) + + identity := schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-one"}) + assert.Nil(t, GetCachedModelSchema(identity), "nothing is cached before the first store") + + StoreModelSchema(identity, schemaWithLabel("from-cache")) + + cached := GetCachedModelSchema(identity) + require.NotNil(t, cached) + require.NotNil(t, cached.ModelSchema) + assert.Equal(t, "from-cache", cached.ModelSchema.Model) +} + +// An upgrade changes the unique identifier, so the old entry is unreachable rather than stale. +func TestModelSchemaCacheIsKeyedOnPluginVersion(t *testing.T) { + setupModelSchemaTest(t) + + credentials := map[string]any{"key": "secret-one"} + StoreModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, credentials), schemaWithLabel("old-version")) + + assert.Nil(t, + GetCachedModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V2, credentials)), + "an upgraded plugin must not read the previous version's schema", + ) +} + +func TestModelSchemaCacheIsKeyedOnCredentials(t *testing.T) { + setupModelSchemaTest(t) + + StoreModelSchema( + schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-one"}), + schemaWithLabel("first-credentials"), + ) + + assert.Nil(t, + GetCachedModelSchema(schemaTestIdentity(SCHEMA_TEST_IDENTIFIER_V1, map[string]any{"key": "secret-two"})), + "a different credential set must not read the first one's schema", + ) +} + +func TestModelSchemaCacheSkipsRemotePlugins(t *testing.T) { + setupModelSchemaTest(t) + + identity := schemaTestIdentity(SCHEMA_TEST_REMOTE, map[string]any{"key": "secret-one"}) + StoreModelSchema(identity, schemaWithLabel("debug-build")) + + assert.Nil(t, GetCachedModelSchema(identity), "a debugging plugin reuses one identifier across builds") +}