feat: add plugin reinstallation functionality and admin API support (#285)

* feat: add plugin reinstallation functionality and admin API support

- Implemented ReinstallToAWSFromPkg method to allow reinstallation of plugins on AWS Lambda, updating function URL and name.
- Added clearServerlessRuntimeCache method to manage serverless runtime cache.
- Enhanced LaunchPlugin to support an ignoreIdempotent flag for forced reinstallation.
- Introduced admin API endpoints for plugin reinstallation, secured with an API key validation middleware.
- Updated configuration to include AdminApiEnabled and AdminApiKey settings.

* refactor: update plugin reinstallation endpoint and improve unauthorized response

- Changed the plugin reinstallation endpoint from "/plugins/reinstall" to "/plugin/serverless/reinstall" for better clarity.
- Modified the unauthorized response in the AdminAPIKey middleware to return a more descriptive JSON message.

Author: Yeuoly

internal/core/plugin_manager/install_to_serverless.go

@@ -34,7 +34,8 @@ func (p *PluginManager) InstallToAWSFromPkg(
 		return nil, err
 	}
 
-	response, err := serverless.LaunchPlugin(originalPackager, decoder, p.serverlessConnectorLaunchTimeout)
+	// serverless.LaunchPlugin will check if the plugin has already been launched, if so, it returns directly
+	response, err := serverless.LaunchPlugin(originalPackager, decoder, p.serverlessConnectorLaunchTimeout, false)
 	if err != nil {
 		return nil, err
 	}
@@ -119,3 +120,120 @@ func (p *PluginManager) InstallToAWSFromPkg(
 
 	return newResponse, nil
 }
+
+/*
+ * Reinstall a plugin to AWS Lambda, update function url and name
+ */
+func (p *PluginManager) ReinstallToAWSFromPkg(
+	originalPackager []byte,
+	decoder decoder.PluginDecoder,
+) (
+	*stream.Stream[PluginInstallResponse], error,
+) {
+	checksum, err := decoder.Checksum()
+	if err != nil {
+		return nil, err
+	}
+	// check valid manifest
+	_, err = decoder.Manifest()
+	if err != nil {
+		return nil, err
+	}
+	uniqueIdentity, err := decoder.UniqueIdentity()
+	if err != nil {
+		return nil, err
+	}
+
+	// check if serverless runtime exists
+	serverlessRuntime, err := db.GetOne[models.ServerlessRuntime](
+		db.Equal("plugin_unique_identifier", uniqueIdentity.String()),
+	)
+	if err == db.ErrDatabaseNotFound {
+		return nil, fmt.Errorf("plugin not exists")
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	response, err := serverless.LaunchPlugin(
+		originalPackager,
+		decoder,
+		p.serverlessConnectorLaunchTimeout,
+		true, // ignoreIdempotent, true means always reinstall
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	newResponse := stream.NewStream[PluginInstallResponse](128)
+	routine.Submit(map[string]string{
+		"module":          "plugin_manager",
+		"function":        "ReinstallToAWSFromPkg",
+		"checksum":        checksum,
+		"unique_identity": uniqueIdentity.String(),
+	}, func() {
+		defer func() {
+			newResponse.Close()
+		}()
+
+		functionUrl := ""
+		functionName := ""
+
+		response.Async(func(r serverless.LaunchFunctionResponse) {
+			if r.Event == serverless.Info {
+				newResponse.Write(PluginInstallResponse{
+					Event: PluginInstallEventInfo,
+					Data:  "Installing...",
+				})
+			} else if r.Event == serverless.Done {
+				if functionUrl == "" || functionName == "" {
+					newResponse.Write(PluginInstallResponse{
+						Event: PluginInstallEventError,
+						Data:  "Internal server error, failed to get lambda url or function name",
+					})
+					return
+				}
+
+				// update serverless runtime
+				serverlessRuntime.FunctionURL = functionUrl
+				serverlessRuntime.FunctionName = functionName
+				err = db.Update(&serverlessRuntime)
+				if err != nil {
+					newResponse.Write(PluginInstallResponse{
+						Event: PluginInstallEventError,
+						Data:  "Failed to update serverless runtime",
+					})
+					return
+				}
+
+				// clear cache
+				err = p.clearServerlessRuntimeCache(uniqueIdentity)
+				if err != nil {
+					newResponse.Write(PluginInstallResponse{
+						Event: PluginInstallEventError,
+						Data:  "Failed to clear serverless runtime cache",
+					})
+					return
+				}
+
+				newResponse.Write(PluginInstallResponse{
+					Event: PluginInstallEventDone,
+					Data:  "Installed",
+				})
+			} else if r.Event == serverless.Error {
+				newResponse.Write(PluginInstallResponse{
+					Event: PluginInstallEventError,
+					Data:  "Internal server error",
+				})
+			} else if r.Event == serverless.FunctionUrl {
+				functionUrl = r.Message
+			} else if r.Event == serverless.Function {
+				functionName = r.Message
+			} else {
+				newResponse.WriteError(fmt.Errorf("unknown event: %s, with message: %s", r.Event, r.Message))
+			}
+		})
+	})
+
+	return newResponse, nil
+}

internal/core/plugin_manager/serverless.go

@@ -93,3 +93,10 @@ func (p *PluginManager) getServerlessPluginRuntimeModel(
 
 	return runtime, nil
 }
+
+func (p *PluginManager) clearServerlessRuntimeCache(
+	identity plugin_entities.PluginUniqueIdentifier,
+) error {
+	_, err := cache.Del(p.getServerlessRuntimeCacheKey(identity))
+	return err
+}

internal/core/plugin_manager/serverless_connector/launch.go

@@ -19,6 +19,7 @@ func LaunchPlugin(
 	originPackage []byte,
 	decoder decoder.PluginDecoder,
 	timeout int, // in seconds
+	ignoreIdempotent bool, // if true, never check if the plugin has launched
 ) (*stream.Stream[LaunchFunctionResponse], error) {
 	checksum, err := decoder.Checksum()
 	if err != nil {
@@ -36,28 +37,30 @@ func LaunchPlugin(
 		return nil, err
 	}
 
-	function, err := FetchFunction(manifest, checksum)
-	if err != nil {
-		if err != ErrFunctionNotFound {
-			return nil, err
+	if !ignoreIdempotent {
+		function, err := FetchFunction(manifest, checksum)
+		if err != nil {
+			if err != ErrFunctionNotFound {
+				return nil, err
+			}
+		} else {
+			// found, return directly
+			response := stream.NewStream[LaunchFunctionResponse](3)
+			response.Write(LaunchFunctionResponse{
+				Event:   FunctionUrl,
+				Message: function.FunctionURL,
+			})
+			response.Write(LaunchFunctionResponse{
+				Event:   Function,
+				Message: function.FunctionName,
+			})
+			response.Write(LaunchFunctionResponse{
+				Event:   Done,
+				Message: "",
+			})
+			response.Close()
+			return response, nil
 		}
-	} else {
-		// found, return directly
-		response := stream.NewStream[LaunchFunctionResponse](3)
-		response.Write(LaunchFunctionResponse{
-			Event:   FunctionUrl,
-			Message: function.FunctionURL,
-		})
-		response.Write(LaunchFunctionResponse{
-			Event:   Function,
-			Message: function.FunctionName,
-		})
-		response.Write(LaunchFunctionResponse{
-			Event:   Done,
-			Message: "",
-		})
-		response.Close()
-		return response, nil
 	}
 
 	response, err := SetupFunction(manifest, checksum, bytes.NewReader(originPackage), timeout)

internal/server/constants/constants.go

@@ -1,8 +1,9 @@
 package constants
 
 const (
-	X_PLUGIN_ID = "X-Plugin-ID"
-	X_API_KEY   = "X-Api-Key"
+	X_PLUGIN_ID     = "X-Plugin-ID"
+	X_API_KEY       = "X-Api-Key"
+	X_ADMIN_API_KEY = "X-Admin-Api-Key"
 
 	CONTEXT_KEY_PLUGIN_INSTALLATION      = "plugin_installation"
 	CONTEXT_KEY_PLUGIN_UNIQUE_IDENTIFIER = "plugin_unique_identifier"

internal/server/controllers/plugins.go

@@ -35,12 +35,12 @@ func UploadPlugin(app *app.Config) gin.HandlerFunc {
 
 		tenantId := c.Param("tenant_id")
 		if tenantId == "" {
-			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("Tenant ID is required")).ToResponse())
+			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("tenant ID is required")).ToResponse())
 			return
 		}
 
 		if difyPkgFileHeader.Size > app.MaxPluginPackageSize {
-			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("File size exceeds the maximum limit")).ToResponse())
+			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("file size exceeds the maximum limit")).ToResponse())
 			return
 		}
 
@@ -67,12 +67,12 @@ func UploadBundle(app *app.Config) gin.HandlerFunc {
 
 		tenantId := c.Param("tenant_id")
 		if tenantId == "" {
-			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("Tenant ID is required")).ToResponse())
+			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("tenant ID is required")).ToResponse())
 			return
 		}
 
 		if difyBundleFileHeader.Size > app.MaxBundlePackageSize {
-			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("File size exceeds the maximum limit")).ToResponse())
+			c.JSON(http.StatusOK, exception.BadRequestError(errors.New("file size exceeds the maximum limit")).ToResponse())
 			return
 		}
 
@@ -140,6 +140,16 @@ func InstallPluginFromIdentifiers(app *app.Config) gin.HandlerFunc {
 	}
 }
 
+func ReinstallPluginFromIdentifier(app *app.Config) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		BindRequest(c, func(request struct {
+			PluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier `json:"plugin_unique_identifier" validate:"required,plugin_unique_identifier"`
+		}) {
+			service.ReinstallPluginFromIdentifier(c, app, request.PluginUniqueIdentifier)
+		})
+	}
+}
+
 func FetchPluginInstallationTasks(c *gin.Context) {
 	BindRequest(c, func(request struct {
 		TenantID string `uri:"tenant_id" validate:"required"`

internal/server/http_server.go

@@ -34,6 +34,17 @@ func (app *App) server(config *app.Config) func() {
 	pluginGroup := engine.Group("/plugin/:tenant_id")
 	pprofGroup := engine.Group("/debug/pprof")
 
+	if config.AdminApiEnabled {
+		if len(config.AdminApiKey) < 10 {
+			log.Panic("length of admin api key must be greater than 10")
+		}
+
+		adminGroup := engine.Group("/admin")
+		adminGroup.Use(app.AdminAPIKey(config.AdminApiKey))
+
+		app.adminGroup(adminGroup, config)
+	}
+
 	if config.SentryEnabled {
 		// setup sentry for all groups
 		sentryGroup := []*gin.RouterGroup{
@@ -169,6 +180,10 @@ func (app *App) pluginManagementGroup(group *gin.RouterGroup, config *app.Config
 	group.GET("/agent_strategy", controllers.GetAgentStrategy)
 }
 
+func (app *App) adminGroup(group *gin.RouterGroup, config *app.Config) {
+	group.POST("/plugin/serverless/reinstall", controllers.ReinstallPluginFromIdentifier(config))
+}
+
 func (app *App) pluginAssetGroup(group *gin.RouterGroup) {
 	group.GET("/:id", controllers.GetAsset)
 }

internal/server/middleware.go

@@ -170,3 +170,14 @@ func (app *App) InitClusterID() gin.HandlerFunc {
 		ctx.Next()
 	}
 }
+
+func (app *App) AdminAPIKey(key string) gin.HandlerFunc {
+	return func(ctx *gin.Context) {
+		if ctx.GetHeader(constants.X_ADMIN_API_KEY) != key {
+			ctx.AbortWithStatusJSON(401, gin.H{"message": "unauthorized"})
+			return
+		}
+
+		ctx.Next()
+	}
+}

internal/service/install_plugin.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/gin-gonic/gin"
 	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_manager"
 	"github.com/langgenius/dify-plugin-daemon/internal/db"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
@@ -335,6 +336,38 @@ func InstallPluginFromIdentifiers(
 	return entities.NewSuccessResponse(response)
 }
 
+/*
+ * Reinstall a plugin from a given identifier, no tenant_id is needed
+ */
+func ReinstallPluginFromIdentifier(
+	ctx *gin.Context,
+	config *app.Config,
+	pluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier,
+) {
+	baseSSEService(func() (*stream.Stream[plugin_manager.PluginInstallResponse], error) {
+		if config.Platform != app.PLATFORM_SERVERLESS {
+			return nil, fmt.Errorf("reinstall is only supported on serverless platform")
+		}
+
+		manager := plugin_manager.Manager()
+		pkgFile, err := manager.GetPackage(pluginUniqueIdentifier)
+		if err != nil {
+			return nil, errors.Join(err, errors.New("failed to get package"))
+		}
+
+		zipDecoder, err := decoder.NewZipPluginDecoder(pkgFile)
+		if err != nil {
+			return nil, errors.Join(err, errors.New("failed to create zip decoder"))
+		}
+		stream, err := manager.ReinstallToAWSFromPkg(pkgFile, zipDecoder)
+		if err != nil {
+			return nil, errors.Join(err, errors.New("failed to reinstall plugin"))
+		}
+
+		return stream, nil
+	}, ctx, 1800)
+}
+
 func UpgradePlugin(
 	config *app.Config,
 	tenant_id string,

internal/types/app/config.go

@@ -2,6 +2,7 @@ package app
 
 import (
 	"fmt"
+
 	"github.com/langgenius/dify-plugin-daemon/internal/oss"
 
 	"github.com/go-playground/validator/v10"
@@ -12,6 +13,10 @@ type Config struct {
 	ServerPort uint16 `envconfig:"SERVER_PORT" validate:"required"`
 	ServerKey  string `envconfig:"SERVER_KEY" validate:"required"`
 
+	// admin api enable
+	AdminApiEnabled bool   `envconfig:"ADMIN_API_ENABLED" default:"false"`
+	AdminApiKey     string `envconfig:"ADMIN_API_KEY"`
+
 	// dify inner api
 	DifyInnerApiURL string `envconfig:"DIFY_INNER_API_URL" validate:"required"`
 	DifyInnerApiKey string `envconfig:"DIFY_INNER_API_KEY" validate:"required"`