feat(cache): cache model schemas in the daemon

.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

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 {

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{},

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

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)
+			},
+		)
+	}
+}

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)
 }

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()
+}