feat(serverless): keep long-running invocations alive via periodic activation

.env.example

@@ -173,6 +173,9 @@ DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_ENABLED=false
 # seconds to wait for the plugin to become ready during the activation preflight;
 # if it is not woken up in time the invocation is treated as failed
 DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_TIMEOUT=60
+# seconds between keepalive calls to the connector's activate endpoint while an
+# invocation is in flight; the connector idle time should be at least 2x this interval
+DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_KEEPALIVE_INTERVAL=60
 
 # python interpreter, if you are using local runtime, you should set this path to your python interpreter path
 # otherwise, it should be /usr/bin/python3

docs/runtime/sri.md

@@ -19,6 +19,7 @@ The daemon is configured using the following environment variables:
 | `MAX_SERVERLESS_REQUEST_BYTES` | Maximum serialized request payload size sent to a serverless plugin runtime (in bytes). Default is 5242880 (5 MB). This limit accounts for Lambda Function URL's 6 MB request size limit, with a safety margin for headers and metadata. |
 | `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_ENABLED` | Whether to run the activation preflight before dispatching an invocation. When enabled, the daemon calls `POST /v1/activation/activate` to wake a scaled-to-zero plugin and waits until it is ready. Default is `false`. |
 | `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_TIMEOUT` | Seconds the daemon waits for the plugin to become ready during the activation preflight. If the plugin is not woken up in time, the invocation is treated as failed. Default is `60`. |
+| `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_KEEPALIVE_INTERVAL` | Seconds between keepalive calls to `POST /v1/activation/activate` while an invocation is in flight, preventing the plugin from being scaled down mid-invocation. Only effective when activation is enabled. The connector idle time should be at least 2x this interval. Keepalive failures are logged and do not fail the invocation. Default is `60`. |
 
 ---
 
@@ -165,6 +166,7 @@ Content-Type: application/json
 **Error Handling**
 
 - The daemon bounds the wait with `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_TIMEOUT`. On timeout or any non-`200` response, the invocation is aborted and reported as a failure.
+- After a successful preflight, the daemon re-calls this endpoint every `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_KEEPALIVE_INTERVAL` seconds for the duration of the invocation to renew the activity lease. Keepalive failures are logged but never abort the invocation.
 
 ---
 

docs/runtime/sri_cn.md

@@ -19,6 +19,7 @@ daemon 通过如下环境变量进行配置：
 | `MAX_SERVERLESS_REQUEST_BYTES` | 发送到 serverless 插件运行时的最大序列化请求负载大小（以字节为单位）。默认值为 5242880（5 MB）。此限制考虑了 Lambda Function URL 的 6 MB 请求大小限制，并为 headers 和元数据留有安全余量。 |
 | `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_ENABLED` | 是否在下发调用前执行激活（activation）预检。启用后，daemon 会调用 `POST /v1/activation/activate` 唤醒被缩容至零的插件，并等待其就绪。默认值为 `false`。 |
 | `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_TIMEOUT` | 激活预检中，daemon 等待插件就绪的超时时间（秒）。若插件未在超时时间内被唤醒，则本次调用按失败处理。默认值为 `60`。 |
+| `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_KEEPALIVE_INTERVAL` | 调用执行期间，daemon 周期性调用 `POST /v1/activation/activate` 保持插件活跃的间隔（秒），防止长时间运行的调用被缩容中断。仅在启用激活时生效。connector 的空闲时间应至少为该间隔的 2 倍。keepalive 失败仅记录日志，不会导致调用失败。默认值为 `60`。 |
 
 ---
 
@@ -165,6 +166,7 @@ Content-Type: application/json
 **错误处理**
 
 - daemon 使用 `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_TIMEOUT` 限制等待时长。一旦超时或收到非 `200` 响应，本次调用将被中断并按失败处理。
+- 预检成功后，daemon 会在调用存续期间每隔 `DIFY_PLUGIN_SERVERLESS_CONNECTOR_ACTIVATION_KEEPALIVE_INTERVAL` 秒再次调用该接口以续期活跃租约。keepalive 失败仅记录日志，不会中断调用。
 
 ---
 

internal/core/serverless_runtime/io.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"net/http"
 	"net/url"
+	"runtime"
 	"strings"
 	"time"
 
@@ -26,6 +27,8 @@ import (
 
 const serverlessErrorResponseLimit = 4 * 1024
 
+var activatePlugin = serverless.Activate
+
 // ServerlessPayloadTooLargeError reports request size metadata without exposing payload contents.
 type ServerlessPayloadTooLargeError struct {
 	Action          access_types.PluginAccessAction
@@ -365,7 +368,7 @@ func (r *ServerlessPluginRuntime) WriteContext(
 		}()
 
 		if r.ActivationEnabled {
-			if err := serverless.Activate(
+			if err := activatePlugin(
 				requestCtx,
 				r.LambdaName,
 				time.Duration(r.ActivationTimeout)*time.Second,
@@ -387,6 +390,8 @@ func (r *ServerlessPluginRuntime) WriteContext(
 				})
 				return
 			}
+			stopKeepalive := r.startActivationKeepalive(requestCtx, sessionId, action)
+			defer stopKeepalive()
 		}
 
 		url += "?action=" + string(action)
@@ -548,3 +553,43 @@ func (r *ServerlessPluginRuntime) WriteContext(
 
 	return nil
 }
+
+func (r *ServerlessPluginRuntime) startActivationKeepalive(ctx context.Context, sessionID string, action access_types.PluginAccessAction) func() {
+	if r.ActivationKeepaliveInterval <= 0 {
+		return func() {}
+	}
+	ctx, cancel := context.WithCancel(ctx)
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		defer func() {
+			if err := recover(); err != nil {
+				buf := make([]byte, 1024*1024)
+				n := runtime.Stack(buf, false)
+				log.Error("panic in serverless activation keepalive", "session_id", sessionID, "action", action, "lambda_name", r.LambdaName, "panic", err, "stack_trace", string(buf[:n]))
+			}
+		}()
+		ticker := time.NewTicker(r.ActivationKeepaliveInterval)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+			}
+			if err := activatePlugin(ctx, r.LambdaName, r.ActivationKeepaliveInterval); err != nil && ctx.Err() == nil {
+				log.Warn(
+					"serverless runtime activation keepalive failed",
+					"session_id", sessionID,
+					"action", action,
+					"lambda_name", r.LambdaName,
+					"error", err,
+				)
+			}
+		}
+	}()
+	return func() {
+		cancel()
+		<-done
+	}
+}

internal/core/serverless_runtime/io_test.go

@@ -1062,3 +1062,235 @@ func (t *trackingReadCloser) Close() error {
 	t.bodiesClosed.Add(1)
 	return t.ReadCloser.Close()
 }
+
+func startWriteWithActivation(
+	t *testing.T,
+	handler http.HandlerFunc,
+	activationEnabled bool,
+	keepaliveInterval time.Duration,
+	activate func(ctx context.Context, instanceID string, timeout time.Duration) error,
+) (chan plugin_entities.SessionMessage, chan struct{}) {
+	t.Helper()
+	routine.InitPool(1)
+
+	originalActivate := activatePlugin
+	activatePlugin = activate
+	t.Cleanup(func() { activatePlugin = originalActivate })
+
+	server := httptest.NewServer(handler)
+	t.Cleanup(server.Close)
+
+	runtime := &ServerlessPluginRuntime{
+		Client:                      server.Client(),
+		LambdaURL:                   server.URL,
+		LambdaName:                  "test-lambda",
+		MaxRequestBytes:             1024,
+		MaxRetryTimes:               1,
+		PluginMaxExecutionTimeout:   10,
+		RuntimeBufferSize:           1024,
+		RuntimeMaxBufferSize:        1024,
+		ActivationEnabled:           activationEnabled,
+		ActivationTimeout:           5,
+		ActivationKeepaliveInterval: keepaliveInterval,
+		listeners:                   mapping.Map[string, *entities.Broadcast[plugin_entities.SessionMessage]]{},
+	}
+
+	listener, err := runtime.Listen("test-session")
+	if err != nil {
+		t.Fatalf("listen failed: %v", err)
+	}
+
+	messages := make(chan plugin_entities.SessionMessage, 8)
+	closed := make(chan struct{})
+	listener.Listen(func(message plugin_entities.SessionMessage) {
+		messages <- message
+	})
+	listener.OnClose(func() {
+		close(closed)
+	})
+
+	if err := runtime.Write(
+		"test-session",
+		access_types.PLUGIN_ACCESS_ACTION_INVOKE_TEXT_EMBEDDING,
+		[]byte("{}"),
+	); err != nil {
+		t.Fatalf("write failed: %v", err)
+	}
+
+	return messages, closed
+}
+
+func collectMessagesAfterClose(t *testing.T, messages chan plugin_entities.SessionMessage, closed chan struct{}) []plugin_entities.SessionMessage {
+	t.Helper()
+	select {
+	case <-closed:
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for listener close")
+	}
+	close(messages)
+	received := make([]plugin_entities.SessionMessage, 0, len(messages))
+	for message := range messages {
+		received = append(received, message)
+	}
+	return received
+}
+
+func slowEndSessionHandler(w http.ResponseWriter, r *http.Request) {
+	time.Sleep(150 * time.Millisecond)
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte(
+		"{\"session_id\":\"test-session\",\"event\":\"session\",\"data\":{\"type\":\"end\",\"data\":{\"reason\":\"done\"}}}\n",
+	))
+	w.(http.Flusher).Flush()
+}
+
+func TestWrite_ActivationKeepaliveRenewsDuringLongInvocation(t *testing.T) {
+	var activateCalls atomic.Int32
+	messages, closed := startWriteWithActivation(
+		t,
+		slowEndSessionHandler,
+		true,
+		20*time.Millisecond,
+		func(ctx context.Context, instanceID string, timeout time.Duration) error {
+			activateCalls.Add(1)
+			return nil
+		},
+	)
+
+	received := collectMessagesAfterClose(t, messages, closed)
+
+	if calls := activateCalls.Load(); calls < 4 {
+		t.Fatalf("expected at least 4 activate calls (1 preflight + keepalives), got %d", calls)
+	}
+	hasEnd := false
+	for _, message := range received {
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_ERROR {
+			t.Fatalf("unexpected error message: %s", message.Data)
+		}
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_END {
+			hasEnd = true
+		}
+	}
+	if !hasEnd {
+		t.Fatal("expected END message")
+	}
+}
+
+func TestWrite_ActivationKeepaliveStopsAfterInvocation(t *testing.T) {
+	var activateCalls atomic.Int32
+	messages, closed := startWriteWithActivation(
+		t,
+		slowEndSessionHandler,
+		true,
+		20*time.Millisecond,
+		func(ctx context.Context, instanceID string, timeout time.Duration) error {
+			activateCalls.Add(1)
+			return nil
+		},
+	)
+
+	collectMessagesAfterClose(t, messages, closed)
+
+	snapshot := activateCalls.Load()
+	time.Sleep(100 * time.Millisecond)
+	if calls := activateCalls.Load(); calls != snapshot {
+		t.Fatalf("expected no activate calls after invocation ended, got %d -> %d", snapshot, calls)
+	}
+}
+
+func TestWrite_ActivationKeepaliveFailureDoesNotAbortInvocation(t *testing.T) {
+	var activateCalls atomic.Int32
+	messages, closed := startWriteWithActivation(
+		t,
+		slowEndSessionHandler,
+		true,
+		20*time.Millisecond,
+		func(ctx context.Context, instanceID string, timeout time.Duration) error {
+			if activateCalls.Add(1) == 1 {
+				return nil
+			}
+			return errors.New("activation keepalive failed")
+		},
+	)
+
+	received := collectMessagesAfterClose(t, messages, closed)
+
+	if calls := activateCalls.Load(); calls < 2 {
+		t.Fatalf("expected keepalive to be attempted, got %d activate calls", calls)
+	}
+	hasEnd := false
+	for _, message := range received {
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_ERROR {
+			t.Fatalf("keepalive failure must not produce an error message: %s", message.Data)
+		}
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_END {
+			hasEnd = true
+		}
+	}
+	if !hasEnd {
+		t.Fatal("expected END message")
+	}
+}
+
+func TestWrite_ActivationDisabledSkipsKeepalive(t *testing.T) {
+	var activateCalls atomic.Int32
+	messages, closed := startWriteWithActivation(
+		t,
+		slowEndSessionHandler,
+		false,
+		20*time.Millisecond,
+		func(ctx context.Context, instanceID string, timeout time.Duration) error {
+			activateCalls.Add(1)
+			return nil
+		},
+	)
+
+	received := collectMessagesAfterClose(t, messages, closed)
+
+	if calls := activateCalls.Load(); calls != 0 {
+		t.Fatalf("expected no activate calls when activation is disabled, got %d", calls)
+	}
+	hasEnd := false
+	for _, message := range received {
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_ERROR {
+			t.Fatalf("unexpected error message: %s", message.Data)
+		}
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_END {
+			hasEnd = true
+		}
+	}
+	if !hasEnd {
+		t.Fatal("expected END message")
+	}
+}
+
+func TestWrite_ActivationKeepalivePanicDoesNotCrashOrHang(t *testing.T) {
+	var activateCalls atomic.Int32
+	messages, closed := startWriteWithActivation(
+		t,
+		slowEndSessionHandler,
+		true,
+		20*time.Millisecond,
+		func(ctx context.Context, instanceID string, timeout time.Duration) error {
+			if activateCalls.Add(1) == 1 {
+				return nil
+			}
+			panic("keepalive panic")
+		},
+	)
+
+	received := collectMessagesAfterClose(t, messages, closed)
+
+	hasEnd := false
+	for _, message := range received {
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_ERROR {
+			t.Fatalf("unexpected error message: %s", message.Data)
+		}
+		if message.Type == plugin_entities.SESSION_MESSAGE_TYPE_END {
+			hasEnd = true
+		}
+	}
+	if !hasEnd {
+		t.Fatal("expected END message")
+	}
+}

internal/core/serverless_runtime/type.go

@@ -38,6 +38,9 @@ type ServerlessPluginRuntime struct {
 	// ActivationTimeout bounds, in seconds, how long the preflight waits for the
 	// plugin to become ready before treating the invocation as failed.
 	ActivationTimeout int
+	// ActivationKeepaliveInterval is how often the daemon re-calls the
+	// connector's activate endpoint while an invocation is in flight.
+	ActivationKeepaliveInterval time.Duration
 
 	RuntimeBufferSize    int
 	RuntimeMaxBufferSize int
@@ -66,16 +69,17 @@ func ConstructServerlessPluginRuntime(
 			MediaTransport: basic_runtime.NewMediaTransport(mediaBucket),
 			InnerChecksum:  serverlessModel.Checksum,
 		},
-		PluginRuntime:             runtimeEntity,
-		LambdaURL:                 serverlessModel.FunctionURL,
-		LambdaName:                serverlessModel.FunctionName,
-		PluginMaxExecutionTimeout: config.PluginMaxExecutionTimeout,
-		MaxRetryTimes:             config.MaxServerlessRetryTimes,
-		MaxRequestBytes:           config.MaxServerlessRequestBytes,
-		ActivationEnabled:         config.DifyPluginServerlessConnectorActivationEnabled,
-		ActivationTimeout:         config.DifyPluginServerlessConnectorActivationTimeout,
-		RuntimeBufferSize:         config.PluginRuntimeBufferSize,
-		RuntimeMaxBufferSize:      config.PluginRuntimeMaxBufferSize,
+		PluginRuntime:               runtimeEntity,
+		LambdaURL:                   serverlessModel.FunctionURL,
+		LambdaName:                  serverlessModel.FunctionName,
+		PluginMaxExecutionTimeout:   config.PluginMaxExecutionTimeout,
+		MaxRetryTimes:               config.MaxServerlessRetryTimes,
+		MaxRequestBytes:             config.MaxServerlessRequestBytes,
+		ActivationEnabled:           config.DifyPluginServerlessConnectorActivationEnabled,
+		ActivationTimeout:           config.DifyPluginServerlessConnectorActivationTimeout,
+		ActivationKeepaliveInterval: time.Duration(config.DifyPluginServerlessConnectorActivationKeepaliveInterval) * time.Second,
+		RuntimeBufferSize:           config.PluginRuntimeBufferSize,
+		RuntimeMaxBufferSize:        config.PluginRuntimeMaxBufferSize,
 
 		// init http client
 		Client: &http.Client{