feat: support plugin poll feature (#744)

fatelei · WH-2099 · web-flow · commit 827d1ce4a2f0 · 2026-05-21T23:45:57.000+08:00
* feat: support plugin poll

* fix: validate polling payloads

---------

Co-authored-by: WH-2099 &lt;wh2099@pm.me&gt;
diff --git a/internal/core/io_tunnel/access_types/access.go b/internal/core/io_tunnel/access_types/access.go
@@ -45,6 +45,8 @@ const (
 	PLUGIN_ACCESS_ACTION_GET_TEXT_EMBEDDING_NUM_TOKENS                      PluginAccessAction = "get_text_embedding_num_tokens"
 	PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS                               PluginAccessAction = "get_ai_model_schemas"
 	PLUGIN_ACCESS_ACTION_GET_LLM_NUM_TOKENS                                 PluginAccessAction = "get_llm_num_tokens"
+	PLUGIN_ACCESS_ACTION_START_POLLING                                      PluginAccessAction = "start_polling"
+	PLUGIN_ACCESS_ACTION_CHECK_POLLING                                      PluginAccessAction = "check_polling"
 	PLUGIN_ACCESS_ACTION_INVOKE_AGENT_STRATEGY                              PluginAccessAction = "invoke_agent_strategy"
 	PLUGIN_ACCESS_ACTION_GET_AUTHORIZATION_URL                              PluginAccessAction = "get_authorization_url"
 	PLUGIN_ACCESS_ACTION_GET_CREDENTIALS                                    PluginAccessAction = "get_credentials"
@@ -81,6 +83,8 @@ func (p PluginAccessAction) IsValid() bool {
 		p == PLUGIN_ACCESS_ACTION_GET_TEXT_EMBEDDING_NUM_TOKENS ||
 		p == PLUGIN_ACCESS_ACTION_GET_AI_MODEL_SCHEMAS ||
 		p == PLUGIN_ACCESS_ACTION_GET_LLM_NUM_TOKENS ||
+		p == PLUGIN_ACCESS_ACTION_START_POLLING ||
+		p == PLUGIN_ACCESS_ACTION_CHECK_POLLING ||
 		p == PLUGIN_ACCESS_ACTION_INVOKE_AGENT_STRATEGY ||
 		p == PLUGIN_ACCESS_ACTION_GET_AUTHORIZATION_URL ||
 		p == PLUGIN_ACCESS_ACTION_GET_CREDENTIALS ||
diff --git a/internal/core/io_tunnel/model_polling.go b/internal/core/io_tunnel/model_polling.go
@@ -0,0 +1,54 @@
+package io_tunnel
+
+import (
+	"errors"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/core/session_manager"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/model_entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/requests"
+)
+
+func StartPolling(
+	session *session_manager.Session,
+	request *requests.RequestStartPolling,
+) (*model_entities.ModelPollingResult, error) {
+	return invokeSingleResponse[requests.RequestStartPolling, model_entities.ModelPollingResult](session, request)
+}
+
+func CheckPolling(
+	session *session_manager.Session,
+	request *requests.RequestCheckPolling,
+) (*model_entities.ModelPollingResult, error) {
+	return invokeSingleResponse[requests.RequestCheckPolling, model_entities.ModelPollingResult](session, request)
+}
+
+func invokeSingleResponse[Req any, Rsp any](
+	session *session_manager.Session,
+	request *Req,
+) (*Rsp, error) {
+	response, err := GenericInvokePlugin[Req, Rsp](session, request, 1)
+	if err != nil {
+		return nil, err
+	}
+	defer response.Close()
+
+	var (
+		got  bool
+		item Rsp
+	)
+
+	for response.Next() {
+		value, readErr := response.Read()
+		if readErr != nil {
+			return nil, readErr
+		}
+		item = value
+		got = true
+	}
+
+	if !got {
+		return nil, errors.New("no polling result received from plugin")
+	}
+
+	return &item, nil
+}
diff --git a/internal/server/controllers/model_polling.go b/internal/server/controllers/model_polling.go
@@ -0,0 +1,28 @@
+package controllers
+
+import (
+	"github.com/gin-gonic/gin"
+	"github.com/langgenius/dify-plugin-daemon/internal/service"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/requests"
+)
+
+func StartPolling() gin.HandlerFunc {
+	type request = plugin_entities.InvokePluginRequest[requests.RequestStartPolling]
+
+	return func(c *gin.Context) {
+		BindPluginDispatchRequest(c, func(itr request) {
+			service.StartPolling(&itr, c)
+		})
+	}
+}
+
+func CheckPolling() gin.HandlerFunc {
+	type request = plugin_entities.InvokePluginRequest[requests.RequestCheckPolling]
+
+	return func(c *gin.Context) {
+		BindPluginDispatchRequest(c, func(itr request) {
+			service.CheckPolling(&itr, c)
+		})
+	}
+}
diff --git a/internal/server/http_server.go b/internal/server/http_server.go
@@ -120,6 +120,8 @@ 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/polling/start", controllers.StartPolling())
+	group.POST("/model/polling/check", controllers.CheckPolling())
 
 	app.setupGeneratedRoutes(group, config)
 }
diff --git a/internal/service/model_polling.go b/internal/service/model_polling.go
@@ -0,0 +1,63 @@
+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/internal/types/exception"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
+	"github.com/langgenius/dify-plugin-daemon/pkg/entities/requests"
+)
+
+func StartPolling(
+	r *plugin_entities.InvokePluginRequest[requests.RequestStartPolling],
+	ctx *gin.Context,
+) {
+	session, err := createSession(
+		r,
+		access_types.PLUGIN_ACCESS_TYPE_MODEL,
+		access_types.PLUGIN_ACCESS_ACTION_START_POLLING,
+		ctx.GetString("cluster_id"),
+		ctx.Request.Context(),
+	)
+	if err != nil {
+		ctx.JSON(500, exception.InternalServerError(err).ToResponse())
+		return
+	}
+	defer session.Close(session_manager.CloseSessionPayload{IgnoreCache: false})
+
+	resp, err := io_tunnel.StartPolling(session, &r.Data)
+	if err != nil {
+		ctx.JSON(500, exception.InvokePluginError(err).ToResponse())
+		return
+	}
+	ctx.JSON(200, entities.NewSuccessResponse(resp))
+}
+
+func CheckPolling(
+	r *plugin_entities.InvokePluginRequest[requests.RequestCheckPolling],
+	ctx *gin.Context,
+) {
+	session, err := createSession(
+		r,
+		access_types.PLUGIN_ACCESS_TYPE_MODEL,
+		access_types.PLUGIN_ACCESS_ACTION_CHECK_POLLING,
+		ctx.GetString("cluster_id"),
+		ctx.Request.Context(),
+	)
+	if err != nil {
+		ctx.JSON(500, exception.InternalServerError(err).ToResponse())
+		return
+	}
+	defer session.Close(session_manager.CloseSessionPayload{IgnoreCache: false})
+
+	resp, err := io_tunnel.CheckPolling(session, &r.Data)
+	if err != nil {
+		ctx.JSON(500, exception.InvokePluginError(err).ToResponse())
+		return
+	}
+
+	ctx.JSON(200, entities.NewSuccessResponse(resp))
+}
diff --git a/pkg/entities/model_entities/polling.go b/pkg/entities/model_entities/polling.go
@@ -0,0 +1,62 @@
+package model_entities
+
+import (
+	"bytes"
+	"encoding/json"
+	"strings"
+
+	"github.com/go-playground/validator/v10"
+	"github.com/langgenius/dify-plugin-daemon/pkg/validators"
+)
+
+type PollingStatus string
+
+const (
+	PollingStatusRunning   PollingStatus = "running"
+	PollingStatusSucceeded PollingStatus = "succeeded"
+	PollingStatusFailed    PollingStatus = "failed"
+)
+
+type ModelPollingResult struct {
+	Status                PollingStatus   `json:"status" validate:"required,oneof=running succeeded failed"`
+	PluginState           map[string]any  `json:"plugin_state,omitempty"`
+	Result                json.RawMessage `json:"result,omitempty"`
+	Error                 string          `json:"error,omitempty"`
+	NextCheckAfterSeconds *int            `json:"next_check_after_seconds,omitempty"`
+	ExpiresAfterSeconds   *int            `json:"expires_after_seconds,omitempty"`
+	MaxAttempts           *int            `json:"max_attempts,omitempty"`
+}
+
+func init() {
+	validators.GlobalEntitiesValidator.RegisterStructValidation(
+		validateModelPollingResult,
+		ModelPollingResult{},
+	)
+}
+
+func validateModelPollingResult(sl validator.StructLevel) {
+	result, ok := sl.Current().Interface().(ModelPollingResult)
+	if !ok {
+		return
+	}
+
+	switch result.Status {
+	case PollingStatusRunning:
+		if len(result.PluginState) == 0 {
+			sl.ReportError(result.PluginState, "plugin_state", "PluginState", "required_for_running", "")
+		}
+	case PollingStatusSucceeded:
+		if isEmptyJSON(result.Result) {
+			sl.ReportError(result.Result, "result", "Result", "required_for_succeeded", "")
+		}
+	case PollingStatusFailed:
+		if strings.TrimSpace(result.Error) == "" {
+			sl.ReportError(result.Error, "error", "Error", "required_for_failed", "")
+		}
+	}
+}
+
+func isEmptyJSON(data json.RawMessage) bool {
+	trimmed := bytes.TrimSpace(data)
+	return len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null"))
+}
diff --git a/pkg/entities/model_entities/polling_test.go b/pkg/entities/model_entities/polling_test.go
@@ -0,0 +1,89 @@
+package model_entities
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/langgenius/dify-plugin-daemon/pkg/validators"
+)
+
+func TestModelPollingResultValidatesStatus(t *testing.T) {
+	nextCheckAfterSeconds := 10
+	result := ModelPollingResult{
+		Status:                PollingStatus("running"),
+		PluginState:           map[string]any{"task_id": "task-1"},
+		NextCheckAfterSeconds: &nextCheckAfterSeconds,
+	}
+
+	if err := validators.GlobalEntitiesValidator.Struct(result); err != nil {
+		t.Fatalf("validate polling result: %v", err)
+	}
+
+	data, err := json.Marshal(result)
+	if err != nil {
+		t.Fatalf("marshal polling result: %v", err)
+	}
+	if !json.Valid(data) {
+		t.Fatalf("invalid json: %s", data)
+	}
+}
+
+func TestModelPollingResultRejectsUnknownStatus(t *testing.T) {
+	result := ModelPollingResult{
+		Status: PollingStatus("pending"),
+	}
+
+	if err := validators.GlobalEntitiesValidator.Struct(result); err == nil {
+		t.Fatal("expected unknown polling status validation error")
+	}
+}
+
+func TestModelPollingResultRequiresRunningState(t *testing.T) {
+	result := ModelPollingResult{
+		Status: PollingStatus("running"),
+	}
+
+	if err := validators.GlobalEntitiesValidator.Struct(result); err == nil {
+		t.Fatal("expected missing plugin_state validation error")
+	}
+}
+
+func TestModelPollingResultRequiresSucceededResult(t *testing.T) {
+	result := ModelPollingResult{
+		Status: PollingStatus("succeeded"),
+		Result: json.RawMessage("null"),
+	}
+
+	if err := validators.GlobalEntitiesValidator.Struct(result); err == nil {
+		t.Fatal("expected missing result validation error")
+	}
+}
+
+func TestModelPollingResultRequiresFailedError(t *testing.T) {
+	result := ModelPollingResult{
+		Status: PollingStatus("failed"),
+		Error:  "  ",
+	}
+
+	if err := validators.GlobalEntitiesValidator.Struct(result); err == nil {
+		t.Fatal("expected missing error validation error")
+	}
+}
+
+func TestModelPollingResultAcceptsTerminalPayloads(t *testing.T) {
+	succeeded := ModelPollingResult{
+		Status: PollingStatus("succeeded"),
+		Result: json.RawMessage(`{"text":"done"}`),
+	}
+	if err := validators.GlobalEntitiesValidator.Struct(succeeded); err != nil {
+		t.Fatalf("validate succeeded result: %v", err)
+	}
+
+	failed := ModelPollingResult{
+		Status: PollingStatus("failed"),
+		Error:  "provider failed",
+	}
+	if err := validators.GlobalEntitiesValidator.Struct(failed); err != nil {
+		t.Fatalf("validate failed result: %v", err)
+	}
+}
diff --git a/pkg/entities/requests/model.go b/pkg/entities/requests/model.go
@@ -21,6 +21,7 @@ type InvokeLLMSchema struct {
 	PromptMessages  []model_entities.PromptMessage     `json:"prompt_messages"  validate:"omitempty"`
 	Tools           []model_entities.PromptMessageTool `json:"tools" validate:"omitempty,dive"`
 	Stop            []string                           `json:"stop" validate:"omitempty"`
+	JSONSchema      map[string]any                     `json:"json_schema" validate:"omitempty"`
 	Stream          bool                               `json:"stream"`
 }
 
@@ -182,3 +183,20 @@ type RequestGetAIModelSchema struct {
 
 	ModelType model_entities.ModelType `json:"model_type"  validate:"required,model_type"`
 }
+
+type RequestStartPolling struct {
+	RequestInvokeLLM
+
+	WorkflowRunID string `json:"workflow_run_id" validate:"required"`
+	NodeID        string `json:"node_id" validate:"required"`
+}
+
+type RequestCheckPolling struct {
+	BaseRequestInvokeModel
+	Credentials
+
+	ModelType     model_entities.ModelType `json:"model_type" validate:"required,model_type,eq=llm"`
+	WorkflowRunID string                   `json:"workflow_run_id" validate:"required"`
+	NodeID        string                   `json:"node_id" validate:"required"`
+	PluginState   map[string]any           `json:"plugin_state" validate:"required,min=1"`
+}
diff --git a/pkg/entities/requests/model_polling_test.go b/pkg/entities/requests/model_polling_test.go

Original file line number	Diff line number	Diff line change
`@@ -120,6 +120,8 @@ func (app App) pluginDispatchGroup(group gin.RouterGroup, config *app.Config)`
`120`	`120`	`group.Use(app.InitClusterID())`
`121`	`121`
`122`	`122`	`group.POST("/agent_strategy/invoke", controllers.InvokeAgentStrategy(config))`
	`123`	`+ group.POST("/model/polling/start", controllers.StartPolling())`
	`124`	`+ group.POST("/model/polling/check", controllers.CheckPolling())`
`123`	`125`
`124`	`126`	`app.setupGeneratedRoutes(group, config)`
`125`	`127`	`}`