dify

dify/api/core/app/layers/pause_state_persist_layer.py

@@ -0,0 +1,133 @@
+from typing import Annotated, Literal, Self, TypeAlias
+
+from pydantic import BaseModel, Field
+from sqlalchemy import Engine
+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
+from core.workflow.graph_events.base import GraphEngineEvent
+from core.workflow.graph_events.graph import GraphRunPausedEvent
+from models.model import AppMode
+from repositories.api_workflow_run_repository import APIWorkflowRunRepository
+from repositories.factory import DifyAPIRepositoryFactory
+
+
+# Wrapper types for `WorkflowAppGenerateEntity` and
+# `AdvancedChatAppGenerateEntity`. These wrappers enable type discrimination
+# and correct reconstruction of the entity field during (de)serialization.
+class _WorkflowGenerateEntityWrapper(BaseModel):
+    type: Literal[AppMode.WORKFLOW] = AppMode.WORKFLOW
+    entity: WorkflowAppGenerateEntity
+
+
+class _AdvancedChatAppGenerateEntityWrapper(BaseModel):
+    type: Literal[AppMode.ADVANCED_CHAT] = AppMode.ADVANCED_CHAT
+    entity: AdvancedChatAppGenerateEntity
+
+
+_GenerateEntityUnion: TypeAlias = Annotated[
+    _WorkflowGenerateEntityWrapper | _AdvancedChatAppGenerateEntityWrapper,
+    Field(discriminator="type"),
+]
+
+
+class WorkflowResumptionContext(BaseModel):
+    """WorkflowResumptionContext captures all state necessary for resumption."""
+
+    version: Literal["1"] = "1"
+
+    # Only workflow / chatflow could be paused.
+    generate_entity: _GenerateEntityUnion
+    serialized_graph_runtime_state: str
+
+    def dumps(self) -> str:
+        return self.model_dump_json()
+
+    @classmethod
+    def loads(cls, value: str) -> Self:
+        return cls.model_validate_json(value)
+
+    def get_generate_entity(self) -> WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity:
+        return self.generate_entity.entity
+
+
+class PauseStatePersistenceLayer(GraphEngineLayer):
+    def __init__(
+        self,
+        session_factory: Engine | sessionmaker[Session],
+        generate_entity: WorkflowAppGenerateEntity | AdvancedChatAppGenerateEntity,
+        state_owner_user_id: str,
+    ):
+        """Create a PauseStatePersistenceLayer.
+
+        The `state_owner_user_id` is used when creating state file for pause.
+        It generally should id of the creator of workflow.
+        """
+        if isinstance(session_factory, Engine):
+            session_factory = sessionmaker(session_factory)
+        self._session_maker = session_factory
+        self._state_owner_user_id = state_owner_user_id
+        self._generate_entity = generate_entity
+
+    def _get_repo(self) -> APIWorkflowRunRepository:
+        return DifyAPIRepositoryFactory.create_api_workflow_run_repository(self._session_maker)
+
+    def on_graph_start(self) -> None:
+        """
+        Called when graph execution starts.
+
+        This is called after the engine has been initialized but before any nodes
+        are executed. Layers can use this to set up resources or log start information.
+        """
+        pass
+
+    def on_event(self, event: GraphEngineEvent) -> None:
+        """
+        Called for every event emitted by the engine.
+
+        This method receives all events generated during graph execution, including:
+        - Graph lifecycle events (start, success, failure)
+        - Node execution events (start, success, failure, retry)
+        - Stream events for response nodes
+        - Container events (iteration, loop)
+
+        Args:
+            event: The event emitted by the engine
+        """
+        if not isinstance(event, GraphRunPausedEvent):
+            return
+
+        assert self.graph_runtime_state is not None
+
+        entity_wrapper: _GenerateEntityUnion
+        if isinstance(self._generate_entity, WorkflowAppGenerateEntity):
+            entity_wrapper = _WorkflowGenerateEntityWrapper(entity=self._generate_entity)
+        else:
+            entity_wrapper = _AdvancedChatAppGenerateEntityWrapper(entity=self._generate_entity)
+
+        state = WorkflowResumptionContext(
+            serialized_graph_runtime_state=self.graph_runtime_state.dumps(),
+            generate_entity=entity_wrapper,
+        )
+
+        workflow_run_id: str | None = self.graph_runtime_state.system_variable.workflow_execution_id
+        assert workflow_run_id is not None
+        repo = self._get_repo()
+        repo.create_workflow_pause(
+            workflow_run_id=workflow_run_id,
+            state_owner_user_id=self._state_owner_user_id,
+            state=state.dumps(),
+        )
+
+    def on_graph_end(self, error: Exception | None) -> None:
+        """
+        Called when graph execution ends.
+
+        This is called after all nodes have been executed or when execution is
+        aborted. Layers can use this to clean up resources or log final state.
+
+        Args:
+            error: The exception that caused execution to fail, or None if successful
+        """
+        pass

dify/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

dify/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)

dify/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