refactor(workflow): prepare inputs by entry node type 🤖🤖🤖 (#42595)

api/controllers/console/app/workflow.py

@@ -53,7 +53,6 @@ from controllers.console.wraps import (
 from controllers.web.error import InvokeRateLimitError as InvokeRateLimitHttpError
 from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
 from core.app.apps.base_app_queue_manager import AppQueueManager
-from core.app.apps.workflow.app_generator import SKIP_PREPARE_USER_INPUTS_KEY
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.file_access import DatabaseFileAccessController
 from core.db.session_factory import session_factory
@@ -1737,15 +1736,12 @@ class DraftWorkflowTriggerRunApi(Resource):
             event = poller.poll()
             if not event:
                 return jsonable_encoder({"status": "waiting", "retry_in": LISTENING_RETRY_IN})
-            workflow_args = dict(event.workflow_args)
-
-            workflow_args[SKIP_PREPARE_USER_INPUTS_KEY] = True
             return helper.compact_generate_response(
                 AppGenerateService.generate(
                     session=session,
                     app_model=app_model,
                     user=current_user,
-                    args=workflow_args,
+                    args=event.workflow_args,
                     invoke_from=InvokeFrom.DEBUGGER,
                     streaming=True,
                     root_node_id=node_id,
@@ -1894,14 +1890,11 @@ class DraftWorkflowTriggerRunAllApi(Resource):
             return jsonable_encoder({"status": "waiting", "retry_in": LISTENING_RETRY_IN})
 
         try:
-            workflow_args = dict(trigger_debug_event.workflow_args)
-
-            workflow_args[SKIP_PREPARE_USER_INPUTS_KEY] = True
             response = AppGenerateService.generate(
                 session=session,
                 app_model=app_model,
                 user=current_user,
-                args=workflow_args,
+                args=trigger_debug_event.workflow_args,
                 invoke_from=InvokeFrom.DEBUGGER,
                 streaming=True,
                 root_node_id=trigger_debug_event.node_id,

api/core/app/apps/workflow/app_generator.py

@@ -41,6 +41,8 @@ from core.helper.trace_id_helper import (
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.repositories import DifyCoreRepositoryFactory
 from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository
+from core.trigger.constants import is_trigger_node_type
+from core.workflow.node_factory import get_default_root_node_id
 from extensions.ext_database import db
 from factories import file_factory
 from graphon.filters import ResponseStreamFilter
@@ -58,8 +60,6 @@ from services.workflow_draft_variable_service import DraftVarLoader, WorkflowDra
 if TYPE_CHECKING:
     from controllers.console.app.workflow import LoopNodeRunPayload
 
-SKIP_PREPARE_USER_INPUTS_KEY = "_skip_prepare_user_inputs"
-
 logger = logging.getLogger(__name__)
 
 
@@ -89,10 +89,6 @@ class WorkflowAppGenerator(BaseAppGenerator):
             return workflow
         return SnippetGenerateService.ensure_start_node_for_worker(workflow, snippet)
 
-    @staticmethod
-    def _should_prepare_user_inputs(args: Mapping[str, Any]) -> bool:
-        return not bool(args.get(SKIP_PREPARE_USER_INPUTS_KEY))
-
     @overload
     def generate(
         self,
@@ -201,9 +197,11 @@ class WorkflowAppGenerator(BaseAppGenerator):
                 **extract_trace_session_id_from_args(args),
             }
             workflow_run_id = str(workflow_run_id or uuid.uuid4())
-            # FIXME (Yeuoly): we need to remove the SKIP_PREPARE_USER_INPUTS_KEY from the args
-            # trigger shouldn't prepare user inputs
-            if self._should_prepare_user_inputs(args):
+            root_node_id = root_node_id or get_default_root_node_id(workflow.graph_dict)
+            root_node_config = workflow.get_node_config_by_id(root_node_id)
+            root_node_type = workflow.get_node_type_from_node_config(root_node_config)
+            # Trigger inputs are already adapted event data, not Start form fields.
+            if not is_trigger_node_type(root_node_type):
                 inputs = self._prepare_user_inputs(
                     user_inputs=inputs,
                     variables=app_config.variables,

api/tasks/async_workflow_tasks.py

@@ -7,7 +7,7 @@ with appropriate retry policies and error handling.
 
 import logging
 from datetime import UTC, datetime
-from typing import Any, NotRequired
+from typing import Any
 
 from celery import shared_task
 from sqlalchemy import select
@@ -46,8 +46,6 @@ logger = logging.getLogger(__name__)
 class WorkflowGeneratorArgsDict(TypedDict):
     inputs: dict[str, Any]
     files: list[Any]
-    _skip_prepare_user_inputs: bool
-    workflow_id: NotRequired[str]
 
 
 @shared_task(queue=AsyncWorkflowQueue.PROFESSIONAL_QUEUE)
@@ -103,7 +101,6 @@ def _build_generator_args(trigger_data: TriggerData) -> WorkflowGeneratorArgsDic
     return {
         "inputs": dict(trigger_data.inputs),
         "files": list(trigger_data.files),
-        "_skip_prepare_user_inputs": True,
     }
 
 
@@ -150,13 +147,9 @@ def _execute_workflow_common(
             # Execute workflow using WorkflowAppGenerator
             generator = WorkflowAppGenerator()
 
-            # Prepare args matching AppGenerateService.generate format
+            # Adapt trigger inputs and files for the generator.
             args = _build_generator_args(trigger_data)
 
-            # If workflow_id was specified, add it to args
-            if trigger_data.workflow_id:
-                args["workflow_id"] = str(trigger_data.workflow_id)
-
             pause_config = PauseStateLayerConfig(
                 session_factory=session_factory.get_session_maker(),
                 state_owner_user_id=workflow.created_by,

api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py

@@ -11,7 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker
 import core.app.apps.workflow.app_generator as app_generator_module
 from core.app.app_config.entities import WorkflowUIBasedAppConfig
 from core.app.apps.draft_variable_saver import DraftVariableSaverFactory
-from core.app.apps.workflow.app_generator import SKIP_PREPARE_USER_INPUTS_KEY, WorkflowAppGenerator
+from core.app.apps.workflow.app_generator import WorkflowAppGenerator
 from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
 from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig, PauseStatePersistenceLayer
 from core.ops.ops_trace_manager import TraceQueueManager
@@ -39,6 +39,7 @@ def _workflow(
         kind=kind,
         version="1",
         created_by="creator",
+        graph={"nodes": [{"id": "start", "data": {"type": "start", "title": "Start", "variables": []}}], "edges": []},
     )
 
 
@@ -123,24 +124,6 @@ def _repositories(
     )
 
 
-def test_should_prepare_user_inputs_defaults_to_true():
-    args = {"inputs": {}}
-
-    assert WorkflowAppGenerator()._should_prepare_user_inputs(args)
-
-
-def test_should_prepare_user_inputs_skips_when_flag_truthy():
-    args = {"inputs": {}, SKIP_PREPARE_USER_INPUTS_KEY: True}
-
-    assert not WorkflowAppGenerator()._should_prepare_user_inputs(args)
-
-
-def test_should_prepare_user_inputs_keeps_validation_when_flag_false():
-    args = {"inputs": {}, SKIP_PREPARE_USER_INPUTS_KEY: False}
-
-    assert WorkflowAppGenerator()._should_prepare_user_inputs(args)
-
-
 def test_ensure_snippet_start_node_in_worker_returns_standard_workflow_without_lookup(
     sqlite_session: Session,
 ) -> None:

api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py

@@ -12,7 +12,7 @@ from sqlalchemy.orm import Session, sessionmaker
 from core.app.app_config.entities import AppAdditionalFeatures, WorkflowUIBasedAppConfig
 from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.apps.workflow import app_generator as app_generator_module
-from core.app.apps.workflow.app_generator import SKIP_PREPARE_USER_INPUTS_KEY, WorkflowAppGenerator
+from core.app.apps.workflow.app_generator import WorkflowAppGenerator
 from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
 from core.ops.ops_trace_manager import TraceQueueManager
 from models.enums import EndUserType
@@ -264,12 +264,6 @@ class TestWorkflowAppGeneratorValidation:
         assert result is injected_workflow
         ensure_start_node.assert_called_once_with(workflow, snippet)
 
-    def test_should_prepare_user_inputs(self):
-        generator = WorkflowAppGenerator()
-
-        assert generator._should_prepare_user_inputs({}) is True
-        assert generator._should_prepare_user_inputs({SKIP_PREPARE_USER_INPUTS_KEY: True}) is False
-
     def test_single_iteration_generate_validates_args(self, sqlite_session: Session):
         generator = WorkflowAppGenerator()
 
@@ -491,79 +485,95 @@ class TestWorkflowAppGeneratorHandleResponse:
 
 
 class TestWorkflowAppGeneratorGenerate:
-    @pytest.mark.usefixtures("sqlite_generator_session")
-    def test_generate_skips_prepare_inputs_when_flag_set(
-        self,
-        monkeypatch: pytest.MonkeyPatch,
-        sqlite_session: Session,
-    ):
+    @pytest.fixture
+    def generation(self, monkeypatch: pytest.MonkeyPatch, sqlite_generator_session: Session):
+        """Keep input preparation real and stop at the execution boundary."""
         generator = WorkflowAppGenerator()
-        app = _persist_app(sqlite_session)
-        workflow = _persist_workflow(sqlite_session)
-        user = _persist_end_user(sqlite_session)
+        app = _persist_app(sqlite_generator_session)
+        workflow = _persist_workflow(sqlite_generator_session)
+        user = _persist_end_user(sqlite_generator_session)
+        monkeypatch.setattr(app_generator_module, "TraceQueueManager", Mock(return_value=Mock(spec=TraceQueueManager)))
+        execute = Mock(return_value={"ok": True})
+        monkeypatch.setattr(generator, "_generate", execute)
+        return generator, app, workflow, user, execute
 
-        app_config = WorkflowUIBasedAppConfig(
-            tenant_id=TENANT_ID,
-            app_id=APP_ID,
-            app_mode=AppMode.WORKFLOW,
-            additional_features=AppAdditionalFeatures(),
-            variables=[],
-            workflow_id=WORKFLOW_ID,
-        )
-        repository_session_makers: list[sessionmaker[Session]] = []
-
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.WorkflowAppConfigManager.get_app_config",
-            lambda app_model, workflow: app_config,
-        )
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.FileUploadConfigManager.convert",
-            lambda features_dict, is_vision=False: None,
-        )
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.file_factory.build_from_mappings",
-            lambda **kwargs: [],
-        )
-        DummyTraceQueueManager = type(
-            "_DummyTraceQueueManager",
-            (TraceQueueManager,),
+    @pytest.mark.parametrize("node_type", ["trigger-webhook", "trigger-schedule", "trigger-plugin"])
+    @pytest.mark.parametrize("root_node_id", [None, "selected-trigger"])
+    @pytest.mark.parametrize("invoke_from", [InvokeFrom.SERVICE_API, InvokeFrom.DEBUGGER])
+    def test_generate_preserves_trigger_inputs(self, generation, node_type, root_node_id, invoke_from):
+        generator, app, workflow, user, execute = generation
+        workflow.graph = json.dumps(
             {
-                "__init__": lambda self, app_id=None, user_id=None: (
-                    setattr(self, "app_id", app_id) or setattr(self, "user_id", user_id)
-                )
-            },
+                "nodes": [
+                    {"id": "first-trigger", "data": {"type": node_type, "title": "First trigger"}},
+                    {"id": "selected-trigger", "data": {"type": node_type, "title": "Selected trigger"}},
+                ],
+                "edges": [],
+            }
         )
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.TraceQueueManager",
-            DummyTraceQueueManager,
-        )
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_execution_repository",
-            lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(),
-        )
-        monkeypatch.setattr(
-            "core.app.apps.workflow.app_generator.DifyCoreRepositoryFactory.create_workflow_node_execution_repository",
-            lambda **kwargs: repository_session_makers.append(kwargs["session_factory"]) or SimpleNamespace(),
-        )
-
-        prepare_inputs = pytest.fail
-        monkeypatch.setattr(generator, "_prepare_user_inputs", lambda **kwargs: prepare_inputs())
-
-        monkeypatch.setattr(generator, "_generate", lambda **kwargs: {"ok": True})
+        inputs = {"event": {"message": "hello", "items": [1, 2]}, "enabled": False}
 
         result = generator.generate(
             app_model=app,
             workflow=workflow,
             user=user,
-            args={"inputs": {}, SKIP_PREPARE_USER_INPUTS_KEY: True},
-            invoke_from=InvokeFrom.WEB_APP,
+            args={"inputs": inputs},
+            invoke_from=invoke_from,
             streaming=False,
-            call_depth=0,
+            root_node_id=root_node_id,
         )
 
         assert result == {"ok": True}
-        assert len(repository_session_makers) == 2
-        assert all(factory.kw["bind"] is sqlite_session.get_bind() for factory in repository_session_makers)
+        assert execute.call_args.kwargs["application_generate_entity"].inputs == inputs
+        assert execute.call_args.kwargs["root_node_id"] == (root_node_id or "first-trigger")
+
+    @pytest.mark.parametrize("root_node_id", [None, "start"])
+    def test_generate_prepares_start_inputs(self, generation, root_node_id):
+        generator, app, workflow, user, execute = generation
+        workflow.graph = json.dumps(
+            {
+                "nodes": [
+                    {
+                        "id": "start",
+                        "data": {
+                            "type": "start",
+                            "title": "Start",
+                            "variables": [
+                                {"variable": "question", "label": "Question", "type": "text-input", "required": True},
+                                {"variable": "count", "label": "Count", "type": "number", "default": "3"},
+                            ],
+                        },
+                    }
+                ],
+                "edges": [],
+            }
+        )
+
+        generator.generate(
+            app_model=app,
+            workflow=workflow,
+            user=user,
+            args={"inputs": {"question": "hello", "unknown": "filtered"}},
+            invoke_from=InvokeFrom.SERVICE_API,
+            streaming=False,
+            root_node_id=root_node_id,
+        )
+
+        assert execute.call_args.kwargs["application_generate_entity"].inputs == {"question": "hello", "count": 3}
+        assert execute.call_args.kwargs["root_node_id"] == "start"
+
+        execute.reset_mock()
+        with pytest.raises(ValueError, match="question is required in input form"):
+            generator.generate(
+                app_model=app,
+                workflow=workflow,
+                user=user,
+                args={"inputs": {}},
+                invoke_from=InvokeFrom.SERVICE_API,
+                streaming=False,
+                root_node_id=root_node_id,
+            )
+        execute.assert_not_called()
 
 
 class TestWorkflowAppGeneratorResume:

api/tests/unit_tests/tasks/test_async_workflow_tasks.py

@@ -1,27 +1,10 @@
 import ast
 import inspect
 
-from core.app.apps.workflow.app_generator import SKIP_PREPARE_USER_INPUTS_KEY
 from services.errors.app import WorkflowNotFoundError
-from services.workflow.entities import WebhookTriggerData
 from tasks import async_workflow_tasks
 
 
-def test_build_generator_args_sets_skip_flag_for_webhook():
-    trigger_data = WebhookTriggerData(
-        app_id="app",
-        tenant_id="tenant",
-        workflow_id="workflow",
-        root_node_id="node",
-        inputs={"webhook_data": {"body": {"foo": "bar"}}},
-    )
-
-    args = async_workflow_tasks._build_generator_args(trigger_data)
-
-    assert args[SKIP_PREPARE_USER_INPUTS_KEY] is True
-    assert args["inputs"]["webhook_data"]["body"]["foo"] == "bar"
-
-
 def test_workflow_not_found_error_message_is_a_single_formatted_string():
     # Regression: previously raised with printf-style positional args
     # ("...%s...", id1, id2), which made Exception store them as a tuple