feat: introduce trigger functionality (#27644)

Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: Stream <Stream_2@qq.com>
Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com>
Co-authored-by: zhsama <torvalds@linux.do>
Co-authored-by: Harry <xh001x@hotmail.com>
Co-authored-by: lyzno1 <yuanyouhuilyz@gmail.com>
Co-authored-by: yessenia <yessenia.contact@gmail.com>
Co-authored-by: hjlarry <hjlarry@163.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: WTW0313 <twwu@dify.ai>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

api/core/app/layers/pause_state_persist_layer.py

@@ -2,7 +2,7 @@ from typing import Annotated, Literal, Self, TypeAlias
 
 from pydantic import BaseModel, Field
 from sqlalchemy import Engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, WorkflowAppGenerateEntity
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
@@ -55,7 +55,7 @@ class WorkflowResumptionContext(BaseModel):
 class PauseStatePersistenceLayer(GraphEngineLayer):
     def __init__(
         self,
-        session_factory: Engine | sessionmaker,
+        session_factory: Engine | sessionmaker[Session],
         generate_entity: WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity,
         state_owner_user_id: str,
     ):
@@ -103,10 +103,8 @@ class PauseStatePersistenceLayer(GraphEngineLayer):
         entity_wrapper: _GenerateEntityUnion
         if isinstance(self._generate_entity, WorkflowAppGenerateEntity):
             entity_wrapper = _WorkflowGenerateEntityWrapper(entity=self._generate_entity)
-        elif isinstance(self._generate_entity, AdvancedChatAppGenerateEntity):
-            entity_wrapper = _AdvancedChatAppGenerateEntityWrapper(entity=self._generate_entity)
         else:
-            raise AssertionError(f"unknown entity type: type={type(self._generate_entity)}")
+            entity_wrapper = _AdvancedChatAppGenerateEntityWrapper(entity=self._generate_entity)
 
         state = WorkflowResumptionContext(
             serialized_graph_runtime_state=self.graph_runtime_state.dumps(),

api/core/app/layers/suspend_layer.py

@@ -0,0 +1,21 @@
+from core.workflow.graph_engine.layers.base import GraphEngineLayer
+from core.workflow.graph_events.base import GraphEngineEvent
+from core.workflow.graph_events.graph import GraphRunPausedEvent
+
+
+class SuspendLayer(GraphEngineLayer):
+    """ """
+
+    def on_graph_start(self):
+        pass
+
+    def on_event(self, event: GraphEngineEvent):
+        """
+        Handle the paused event, stash runtime state into storage and wait for resume.
+        """
+        if isinstance(event, GraphRunPausedEvent):
+            pass
+
+    def on_graph_end(self, error: Exception | None):
+        """ """
+        pass

api/core/app/layers/timeslice_layer.py

@@ -0,0 +1,88 @@
+import logging
+import uuid
+from typing import ClassVar
+
+from apscheduler.schedulers.background import BackgroundScheduler  # type: ignore
+
+from core.workflow.graph_engine.entities.commands import CommandType, GraphEngineCommand
+from core.workflow.graph_engine.layers.base import GraphEngineLayer
+from core.workflow.graph_events.base import GraphEngineEvent
+from services.workflow.entities import WorkflowScheduleCFSPlanEntity
+from services.workflow.scheduler import CFSPlanScheduler, SchedulerCommand
+
+logger = logging.getLogger(__name__)
+
+
+class TimeSliceLayer(GraphEngineLayer):
+    """
+    CFS plan scheduler to control the timeslice of the workflow.
+    """
+
+    scheduler: ClassVar[BackgroundScheduler] = BackgroundScheduler()
+
+    def __init__(self, cfs_plan_scheduler: CFSPlanScheduler) -> None:
+        """
+        CFS plan scheduler allows to control the timeslice of the workflow.
+        """
+
+        if not TimeSliceLayer.scheduler.running:
+            TimeSliceLayer.scheduler.start()
+
+        super().__init__()
+        self.cfs_plan_scheduler = cfs_plan_scheduler
+        self.stopped = False
+        self.schedule_id = ""
+
+    def _checker_job(self, schedule_id: str):
+        """
+        Check if the workflow need to be suspended.
+        """
+        try:
+            if self.stopped:
+                self.scheduler.remove_job(schedule_id)
+                return
+
+            if self.cfs_plan_scheduler.can_schedule() == SchedulerCommand.RESOURCE_LIMIT_REACHED:
+                # remove the job
+                self.scheduler.remove_job(schedule_id)
+
+                if not self.command_channel:
+                    logger.exception("No command channel to stop the workflow")
+                    return
+
+                # send command to pause the workflow
+                self.command_channel.send_command(
+                    GraphEngineCommand(
+                        command_type=CommandType.PAUSE,
+                        payload={
+                            "reason": SchedulerCommand.RESOURCE_LIMIT_REACHED,
+                        },
+                    )
+                )
+
+        except Exception:
+            logger.exception("scheduler error during check if the workflow need to be suspended")
+
+    def on_graph_start(self):
+        """
+        Start timer to check if the workflow need to be suspended.
+        """
+
+        if self.cfs_plan_scheduler.plan.schedule_strategy == WorkflowScheduleCFSPlanEntity.Strategy.TimeSlice:
+            self.schedule_id = uuid.uuid4().hex
+
+            self.scheduler.add_job(
+                lambda: self._checker_job(self.schedule_id),
+                "interval",
+                seconds=self.cfs_plan_scheduler.plan.granularity,
+                id=self.schedule_id,
+            )
+
+    def on_event(self, event: GraphEngineEvent):
+        pass
+
+    def on_graph_end(self, error: Exception | None) -> None:
+        self.stopped = True
+        # remove the scheduler
+        if self.schedule_id:
+            self.scheduler.remove_job(self.schedule_id)

api/core/app/layers/trigger_post_layer.py

@@ -0,0 +1,88 @@
+import logging
+from datetime import UTC, datetime
+from typing import Any, ClassVar
+
+from pydantic import TypeAdapter
+from sqlalchemy.orm import Session, sessionmaker
+
+from core.workflow.graph_engine.layers.base import GraphEngineLayer
+from core.workflow.graph_events.base import GraphEngineEvent
+from core.workflow.graph_events.graph import GraphRunFailedEvent, GraphRunPausedEvent, GraphRunSucceededEvent
+from models.enums import WorkflowTriggerStatus
+from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
+from tasks.workflow_cfs_scheduler.cfs_scheduler import AsyncWorkflowCFSPlanEntity
+
+logger = logging.getLogger(__name__)
+
+
+class TriggerPostLayer(GraphEngineLayer):
+    """
+    Trigger post layer.
+    """
+
+    _STATUS_MAP: ClassVar[dict[type[GraphEngineEvent], WorkflowTriggerStatus]] = {
+        GraphRunSucceededEvent: WorkflowTriggerStatus.SUCCEEDED,
+        GraphRunFailedEvent: WorkflowTriggerStatus.FAILED,
+        GraphRunPausedEvent: WorkflowTriggerStatus.PAUSED,
+    }
+
+    def __init__(
+        self,
+        cfs_plan_scheduler_entity: AsyncWorkflowCFSPlanEntity,
+        start_time: datetime,
+        trigger_log_id: str,
+        session_maker: sessionmaker[Session],
+    ):
+        self.trigger_log_id = trigger_log_id
+        self.start_time = start_time
+        self.cfs_plan_scheduler_entity = cfs_plan_scheduler_entity
+        self.session_maker = session_maker
+
+    def on_graph_start(self):
+        pass
+
+    def on_event(self, event: GraphEngineEvent):
+        """
+        Update trigger log with success or failure.
+        """
+        if isinstance(event, tuple(self._STATUS_MAP.keys())):
+            with self.session_maker() as session:
+                repo = SQLAlchemyWorkflowTriggerLogRepository(session)
+                trigger_log = repo.get_by_id(self.trigger_log_id)
+                if not trigger_log:
+                    logger.exception("Trigger log not found: %s", self.trigger_log_id)
+                    return
+
+                # Calculate elapsed time
+                elapsed_time = (datetime.now(UTC) - self.start_time).total_seconds()
+
+                # Extract relevant data from result
+                if not self.graph_runtime_state:
+                    logger.exception("Graph runtime state is not set")
+                    return
+
+                outputs = self.graph_runtime_state.outputs
+
+                # BASICLY, workflow_execution_id is the same as workflow_run_id
+                workflow_run_id = self.graph_runtime_state.system_variable.workflow_execution_id
+                assert workflow_run_id, "Workflow run id is not set"
+
+                total_tokens = self.graph_runtime_state.total_tokens
+
+                # Update trigger log with success
+                trigger_log.status = self._STATUS_MAP[type(event)]
+                trigger_log.workflow_run_id = workflow_run_id
+                trigger_log.outputs = TypeAdapter(dict[str, Any]).dump_json(outputs).decode()
+
+                if trigger_log.elapsed_time is None:
+                    trigger_log.elapsed_time = elapsed_time
+                else:
+                    trigger_log.elapsed_time += elapsed_time
+
+                trigger_log.total_tokens = total_tokens
+                trigger_log.finished_at = datetime.now(UTC)
+                repo.update(trigger_log)
+                session.commit()
+
+    def on_graph_end(self, error: Exception | None) -> None:
+        pass