fix: fix plugin upgrade use context (#639)

Author: fatelei

.env.example

@@ -203,6 +203,9 @@ PLUGIN_STDIO_MAX_BUFFER_SIZE=5242880
 PLUGIN_RUNTIME_BUFFER_SIZE=1024
 PLUGIN_RUNTIME_MAX_BUFFER_SIZE=5242880
 
+# plugin install/upgrade timeout in minutes
+PLUGIN_INSTALL_TIMEOUT=15
+
 # dify backwards invocation write timeout in milliseconds
 DIFY_BACKWARDS_INVOCATION_WRITE_TIMEOUT=5000
 # dify backwards invocation read timeout in milliseconds

internal/server/controllers/plugins.go

@@ -112,7 +112,6 @@ func UpgradePlugin(app *app.Config) gin.HandlerFunc {
 			}
 
 			c.JSON(http.StatusOK, service.UpgradePlugin(
-				c.Request.Context(),
 				app,
 				request.TenantID,
 				request.Source,

internal/service/install_plugin.go

@@ -118,7 +118,7 @@ func InstallMultiplePluginsToTenant(
 	for _, job := range jobs {
 		jobCopy := job
 		// create a detached context for async task to avoid http request cancellation
-		taskCtx, taskCancel := context.WithTimeout(context.Background(), 15*time.Minute)
+		taskCtx, taskCancel := context.WithTimeout(context.Background(), time.Duration(config.PluginInstallTimeout)*time.Minute)
 		// start a new goroutine to install the plugin
 		routine.Submit(routinepkg.Labels{
 			routinepkg.RoutineLabelKeyModule: "service",
@@ -197,7 +197,6 @@ func ReinstallPluginFromIdentifier(
  * Upgrade a plugin between 2 identifiers
  */
 func UpgradePlugin(
-	ctx context.Context,
 	config *app.Config,
 	tenantId string,
 	source string,
@@ -299,8 +298,10 @@ func UpgradePlugin(
 		routinepkg.RoutineLabelKeyModule: "service",
 		routinepkg.RoutineLabelKeyMethod: "UpgradePlugin",
 	}, func() {
+		taskCtx, taskCancel := context.WithTimeout(context.Background(), time.Duration(config.PluginInstallTimeout)*time.Minute)
+		defer taskCancel()
 		tasks.ProcessUpgradeJob(
-			ctx,
+			taskCtx,
 			manager,
 			tenants,
 			runtimeType,

internal/service/install_plugin_test.go

@@ -0,0 +1,88 @@
+package service
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
+)
+
+func TestUpgradePlugin(t *testing.T) {
+	originalIdentifier, err := plugin_entities.NewPluginUniqueIdentifier("author/test-plugin:1.0.0@abcdef1234567890abcdef1234567890ab")
+	if err != nil {
+		t.Fatalf("failed to create original plugin unique identifier: %v", err)
+	}
+
+	newIdentifier, err := plugin_entities.NewPluginUniqueIdentifier("author/test-plugin:2.0.0@1234567890abcdef1234567890abcdef12")
+	if err != nil {
+		t.Fatalf("failed to create new plugin unique identifier: %v", err)
+	}
+
+	config := &app.Config{
+		PluginInstallTimeout: 15,
+	}
+
+	tests := []struct {
+		name                             string
+		tenantId                         string
+		source                           string
+		meta                             map[string]any
+		originalPluginUniqueIdentifier   plugin_entities.PluginUniqueIdentifier
+		newPluginUniqueIdentifier        plugin_entities.PluginUniqueIdentifier
+		wantSuccess                      bool
+		wantAllInstalled                 bool
+		wantTaskIDEmpty                  bool
+	}{
+		{
+			name:                           "same plugin identifiers",
+			tenantId:                       "tenant-123",
+			source:                         "test",
+			meta:                           map[string]any{},
+			originalPluginUniqueIdentifier:  originalIdentifier,
+			newPluginUniqueIdentifier:       originalIdentifier,
+			wantSuccess:                    false,
+		},
+		{
+			name:                           "different plugin identifiers",
+			tenantId:                       "tenant-123",
+			source:                         "test",
+			meta:                           map[string]any{},
+			originalPluginUniqueIdentifier:  originalIdentifier,
+			newPluginUniqueIdentifier:       newIdentifier,
+			wantSuccess:                    false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			response := UpgradePlugin(
+				config,
+				tt.tenantId,
+				tt.source,
+				tt.meta,
+				tt.originalPluginUniqueIdentifier,
+				tt.newPluginUniqueIdentifier,
+			)
+
+			isSuccess := response.Code == 0
+			if isSuccess != tt.wantSuccess {
+				t.Errorf("UpgradePlugin() success = %v, want %v", isSuccess, tt.wantSuccess)
+			}
+
+			if isSuccess {
+				var result InstallPluginResponse
+				dataBytes, err := json.Marshal(response.Data)
+				if err == nil {
+					_ = json.Unmarshal(dataBytes, &result)
+					if tt.wantAllInstalled && !result.AllInstalled {
+						t.Errorf("UpgradePlugin() AllInstalled = %v, want %v", result.AllInstalled, tt.wantAllInstalled)
+					}
+					if tt.wantTaskIDEmpty && result.TaskID != "" {
+						t.Errorf("UpgradePlugin() TaskID = %v, want empty", result.TaskID)
+					}
+				}
+			}
+		})
+	}
+}

internal/types/app/config.go

@@ -105,6 +105,9 @@ type Config struct {
 	// request timeout
 	PluginMaxExecutionTimeout int `envconfig:"PLUGIN_MAX_EXECUTION_TIMEOUT" validate:"required"`
 
+	// plugin install timeout in minutes
+	PluginInstallTimeout int `envconfig:"PLUGIN_INSTALL_TIMEOUT" default:"15"`
+
 	// local launching max concurrent
 	PluginLocalLaunchingConcurrent int `envconfig:"PLUGIN_LOCAL_LAUNCHING_CONCURRENT" validate:"required"`