fix: fix some panic and fix test failed (#715)

Author: fatelei

.github/workflows/benchmark.yml

@@ -34,7 +34,7 @@ jobs:
         env:
           UV_PATH: ${{ env.UV_PATH }}
         run: |
-          go test -run=^$ -bench=. -benchmem -benchtime=30s ./... > result.txt
+          go test -run=^$ -bench=. -benchmem -benchtime=30s ./...
 
       - name: Remove License
         run: git stash

internal/core/control_panel/server_debugger.go

@@ -50,6 +50,10 @@ func (c *ControlPanel) onDebuggingRuntimeConnected(
 func (c *ControlPanel) onDebuggingRuntimeDisconnected(
 	rpr *debugging_runtime.RemotePluginRuntime,
 ) {
+	if !rpr.Initialized() {
+		return
+	}
+
 	// handle plugin disconnecting
 	pluginIdentifier, err := rpr.Identity()
 	if err != nil {

internal/core/control_panel/server_debugger_callbacks_test.go

@@ -1,90 +1,113 @@
 package controlpanel
 
 import (
-    "reflect"
-    "testing"
-    "unsafe"
-
-    "github.com/google/uuid"
-    "github.com/stretchr/testify/assert"
-
-    "github.com/langgenius/dify-plugin-daemon/internal/core/debugging_runtime"
-    "github.com/langgenius/dify-plugin-daemon/internal/core/local_runtime"
-    "github.com/langgenius/dify-plugin-daemon/internal/types/app"
-    "github.com/langgenius/dify-plugin-daemon/pkg/entities/manifest_entities"
-    "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
+	"reflect"
+	"testing"
+	"unsafe"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/core/debugging_runtime"
+	"github.com/langgenius/dify-plugin-daemon/internal/core/local_runtime"
+	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/manifest_entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
 )
 
 type mockNotifier struct {
-    connected    bool
-    disconnected bool
+	connected    bool
+	disconnected bool
 }
 
-func (m *mockNotifier) OnLocalRuntimeStarting(plugin_entities.PluginUniqueIdentifier) {}
-func (m *mockNotifier) OnLocalRuntimeReady(*local_runtime.LocalPluginRuntime) {}
+func (m *mockNotifier) OnLocalRuntimeStarting(plugin_entities.PluginUniqueIdentifier)           {}
+func (m *mockNotifier) OnLocalRuntimeReady(*local_runtime.LocalPluginRuntime)                   {}
 func (m *mockNotifier) OnLocalRuntimeStartFailed(plugin_entities.PluginUniqueIdentifier, error) {}
-func (m *mockNotifier) OnLocalRuntimeStop(*local_runtime.LocalPluginRuntime) {}
-func (m *mockNotifier) OnLocalRuntimeStopped(*local_runtime.LocalPluginRuntime) {}
-func (m *mockNotifier) OnLocalRuntimeScaleUp(*local_runtime.LocalPluginRuntime, int32) {}
-func (m *mockNotifier) OnLocalRuntimeScaleDown(*local_runtime.LocalPluginRuntime, int32) {}
-func (m *mockNotifier) OnLocalRuntimeInstanceLog(*local_runtime.LocalPluginRuntime, *local_runtime.PluginInstance, plugin_entities.PluginLogEvent) {}
-func (m *mockNotifier) OnDebuggingRuntimeConnected(r *debugging_runtime.RemotePluginRuntime)  { m.connected = true }
-func (m *mockNotifier) OnDebuggingRuntimeDisconnected(r *debugging_runtime.RemotePluginRuntime) { m.disconnected = true }
+func (m *mockNotifier) OnLocalRuntimeStop(*local_runtime.LocalPluginRuntime)                    {}
+func (m *mockNotifier) OnLocalRuntimeStopped(*local_runtime.LocalPluginRuntime)                 {}
+func (m *mockNotifier) OnLocalRuntimeScaleUp(*local_runtime.LocalPluginRuntime, int32)          {}
+func (m *mockNotifier) OnLocalRuntimeScaleDown(*local_runtime.LocalPluginRuntime, int32)        {}
+func (m *mockNotifier) OnLocalRuntimeInstanceLog(*local_runtime.LocalPluginRuntime, *local_runtime.PluginInstance, plugin_entities.PluginLogEvent) {
+}
+func (m *mockNotifier) OnDebuggingRuntimeConnected(r *debugging_runtime.RemotePluginRuntime) {
+	m.connected = true
+}
+func (m *mockNotifier) OnDebuggingRuntimeDisconnected(r *debugging_runtime.RemotePluginRuntime) {
+	m.disconnected = true
+}
 
 // setPrivateString sets an unexported string field on a struct value via unsafe reflection.
 func setPrivateString(target any, field string, value string) {
-    rv := reflect.ValueOf(target).Elem()
-    f := rv.FieldByName(field)
-    reflect.NewAt(f.Type(), unsafe.Pointer(f.UnsafeAddr())).Elem().SetString(value)
+	rv := reflect.ValueOf(target).Elem()
+	f := rv.FieldByName(field)
+	reflect.NewAt(f.Type(), unsafe.Pointer(f.UnsafeAddr())).Elem().SetString(value)
 }
 
 func newFakeRemoteRuntime(t *testing.T, name, version string) *debugging_runtime.RemotePluginRuntime {
-    t.Helper()
-
-    r := &debugging_runtime.RemotePluginRuntime{
-        PluginRuntime: plugin_entities.PluginRuntime{
-            Config: plugin_entities.PluginDeclaration{
-                PluginDeclarationWithoutAdvancedFields: plugin_entities.PluginDeclarationWithoutAdvancedFields{
-                    // Author is overwritten with tenantId inside Identity()
-                    Name:    name,
-                    Version: manifest_entities.Version(version),
-                },
-            },
-        },
-    }
-
-    // Provide required tenantId and checksum so Identity() succeeds
-    setPrivateString(r, "tenantId", uuid.New().String())
-    setPrivateString(r, "checksum", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef")
-
-    return r
+	t.Helper()
+
+	r := &debugging_runtime.RemotePluginRuntime{
+		PluginRuntime: plugin_entities.PluginRuntime{
+			Config: plugin_entities.PluginDeclaration{
+				PluginDeclarationWithoutAdvancedFields: plugin_entities.PluginDeclarationWithoutAdvancedFields{
+					// Author is overwritten with tenantId inside Identity()
+					Name:    name,
+					Version: manifest_entities.Version(version),
+				},
+			},
+		},
+	}
+
+	// Provide required tenantId and checksum so Identity() succeeds
+	setPrivateString(r, "tenantId", uuid.New().String())
+	setPrivateString(r, "checksum", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef")
+	reflect.NewAt(
+		reflect.ValueOf(r).Elem().FieldByName("initialized").Type(),
+		unsafe.Pointer(reflect.ValueOf(r).Elem().FieldByName("initialized").UnsafeAddr()),
+	).Elem().SetBool(true)
+
+	return r
 }
 
 func TestOnDebuggingRuntimeConnectedAndDisconnected(t *testing.T) {
-    cp := NewControlPanel(&app.Config{PluginLocalLaunchingConcurrent: 1}, nil, nil, nil, nil)
+	cp := NewControlPanel(&app.Config{PluginLocalLaunchingConcurrent: 1}, nil, nil, nil, nil)
 
-    // Attach mock notifier to observe callbacks
-    mn := &mockNotifier{}
-    cp.AddNotifier(mn)
+	// Attach mock notifier to observe callbacks
+	mn := &mockNotifier{}
+	cp.AddNotifier(mn)
 
-    r := newFakeRemoteRuntime(t, "conn_test", "1.2.3")
+	r := newFakeRemoteRuntime(t, "conn_test", "1.2.3")
 
-    // Connected path
-    err := cp.onDebuggingRuntimeConnected(r)
-    assert.NoError(t, err)
+	// Connected path
+	err := cp.onDebuggingRuntimeConnected(r)
+	assert.NoError(t, err)
 
-    id, err := r.Identity()
-    assert.NoError(t, err)
+	id, err := r.Identity()
+	assert.NoError(t, err)
 
-    // Stored in runtime map and notifier called
-    _, ok := cp.debuggingPluginRuntime.Load(id)
-    assert.True(t, ok, "runtime should be stored on connect")
-    assert.True(t, mn.connected, "connected notifier should be triggered")
+	// Stored in runtime map and notifier called
+	_, ok := cp.debuggingPluginRuntime.Load(id)
+	assert.True(t, ok, "runtime should be stored on connect")
+	assert.True(t, mn.connected, "connected notifier should be triggered")
 
-    // Disconnected path
-    cp.onDebuggingRuntimeDisconnected(r)
+	// Disconnected path
+	cp.onDebuggingRuntimeDisconnected(r)
+
+	_, ok = cp.debuggingPluginRuntime.Load(id)
+	assert.False(t, ok, "runtime should be removed on disconnect")
+	assert.True(t, mn.disconnected, "disconnected notifier should be triggered")
+}
 
-    _, ok = cp.debuggingPluginRuntime.Load(id)
-    assert.False(t, ok, "runtime should be removed on disconnect")
-    assert.True(t, mn.disconnected, "disconnected notifier should be triggered")
-}
+func TestOnDebuggingRuntimeDisconnected_IgnoresUninitializedRuntime(t *testing.T) {
+	cp := NewControlPanel(&app.Config{PluginLocalLaunchingConcurrent: 1}, nil, nil, nil, nil)
+
+	mn := &mockNotifier{}
+	cp.AddNotifier(mn)
+
+	r := &debugging_runtime.RemotePluginRuntime{}
+
+	assert.NotPanics(t, func() {
+		cp.onDebuggingRuntimeDisconnected(r)
+	})
+	assert.False(t, mn.disconnected, "disconnected notifier should not be triggered for uninitialized runtime")
+}

internal/core/debugging_runtime/type.go

@@ -156,6 +156,10 @@ func (r *RemotePluginRuntime) SetInstallationId(installationId string) {
 	r.installationId = installationId
 }
 
+func (r *RemotePluginRuntime) Initialized() bool {
+	return r.initialized
+}
+
 func (r *RemotePluginRuntime) Identity() (plugin_entities.PluginUniqueIdentifier, error) {
 	// FIXME: it's a little bit tricky that replace author with current tenant_id
 	// just as a flag to identify debugging plugin

internal/core/io_tunnel/backwards_invocation/transaction/serverless_handler.go

@@ -1,8 +1,10 @@
 package transaction
 
 import (
+	"fmt"
 	"io"
 	"net/http"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -28,17 +30,46 @@ func NewServerlessTransactionHandler(maxTimeout time.Duration) *ServerlessTransa
 type serverlessTransactionWriteCloser struct {
 	done   chan bool
 	closed int32
+	mu     sync.Mutex
 
 	writer func([]byte) (int, error)
 	flush  func()
 }
 
-func (a *serverlessTransactionWriteCloser) Write(data []byte) (int, error) {
-	return a.writer(data)
+func (w *serverlessTransactionWriteCloser) Write(data []byte) (n int, err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if atomic.LoadInt32(&w.closed) != 0 {
+		return 0, io.ErrClosedPipe
+	}
+
+	defer func() {
+		if recovered := recover(); recovered != nil {
+			n = 0
+			err = fmt.Errorf("serverless transaction write panic: %v", recovered)
+			_ = w.Close()
+		}
+	}()
+
+	return w.writer(data)
 }
 
-func (a *serverlessTransactionWriteCloser) Flush() {
-	a.flush()
+func (w *serverlessTransactionWriteCloser) Flush() {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	if atomic.LoadInt32(&w.closed) != 0 {
+		return
+	}
+
+	defer func() {
+		if recover() != nil {
+			_ = w.Close()
+		}
+	}()
+
+	w.flush()
 }
 
 func (w *serverlessTransactionWriteCloser) Close() error {
@@ -64,9 +95,6 @@ func (h *ServerlessTransactionHandler) Handle(ctx *gin.Context, sessionId string
 		return
 	}
 
-	ctx.Writer.WriteHeader(http.StatusOK)
-	ctx.Writer.Header().Set("Content-Type", "text/event-stream")
-
 	plugin_entities.ParsePluginUniversalEvent(
 		bytes,
 		"",
@@ -121,15 +149,17 @@ func (h *ServerlessTransactionHandler) Handle(ctx *gin.Context, sessionId string
 				return
 			}
 
+			ctx.Writer.Header().Set("Content-Type", "text/event-stream")
+
 			// replace trace context, propagate it to gin
 			ctxRequestContext := ctx.Request.Context()
 			ctxRequestContext = log.WithTrace(ctxRequestContext, session.TraceContext)
 			ctxRequestContext = log.WithIdentity(ctxRequestContext, session.IdentityContext)
 			ctx.Request = ctx.Request.WithContext(ctxRequestContext)
 
 			// bind the backwards invocation
-			plugin_manager := plugin_manager.Manager()
-			session.BindBackwardsInvocation(plugin_manager.BackwardsInvocation())
+			pluginManager := plugin_manager.Manager()
+			session.BindBackwardsInvocation(pluginManager.BackwardsInvocation())
 
 			serverlessResponseWriter := NewServerlessTransactionWriter(session, writer)
 
@@ -140,8 +170,10 @@ func (h *ServerlessTransactionHandler) Handle(ctx *gin.Context, sessionId string
 				serverlessResponseWriter,
 				sessionMessage.Data,
 			); err != nil {
-				ctx.Writer.WriteHeader(http.StatusInternalServerError)
-				ctx.Writer.Write([]byte("failed to parse request"))
+				if !ctx.Writer.Written() {
+					ctx.Writer.WriteHeader(http.StatusInternalServerError)
+					ctx.Writer.Write([]byte("failed to parse request"))
+				}
 				writer.Close()
 			}
 		},