Merge branch 'main' into feat/pull-a-variable

api/core/workflow/graph_engine/graph_engine.py

@@ -7,15 +7,13 @@ Domain-Driven Design principles for improved maintainability and testability.
 
 from __future__ import annotations
 
-import contextvars
 import logging
 import queue
 import threading
 from collections.abc import Generator
 from typing import TYPE_CHECKING, cast, final
 
-from flask import Flask, current_app
-
+from core.workflow.context import capture_current_context
 from core.workflow.enums import NodeExecutionType
 from core.workflow.graph import Graph
 from core.workflow.graph_events import (
@@ -159,17 +157,8 @@ class GraphEngine:
         self._layers: list[GraphEngineLayer] = []
 
         # === Worker Pool Setup ===
-        # Capture Flask app context for worker threads
-        flask_app: Flask | None = None
-        try:
-            app = current_app._get_current_object()  # type: ignore
-            if isinstance(app, Flask):
-                flask_app = app
-        except RuntimeError:
-            pass
-
-        # Capture context variables for worker threads
-        context_vars = contextvars.copy_context()
+        # Capture execution context for worker threads
+        execution_context = capture_current_context()
 
         # Create worker pool for parallel node execution
         self._worker_pool = WorkerPool(
@@ -177,8 +166,7 @@ class GraphEngine:
             event_queue=self._event_queue,
             graph=self._graph,
             layers=self._layers,
-            flask_app=flask_app,
-            context_vars=context_vars,
+            execution_context=execution_context,
             min_workers=self._min_workers,
             max_workers=self._max_workers,
             scale_up_threshold=self._scale_up_threshold,

api/core/workflow/graph_engine/worker.py

@@ -5,26 +5,27 @@ Workers pull node IDs from the ready_queue, execute nodes, and push events
 to the event_queue for the dispatcher to process.
 """
 
-import contextvars
 import queue
 import threading
 import time
 from collections.abc import Sequence
 from datetime import datetime
-from typing import final
+from typing import TYPE_CHECKING, final
 from uuid import uuid4
 
-from flask import Flask
 from typing_extensions import override
 
+from core.workflow.context import IExecutionContext
 from core.workflow.graph import Graph
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
 from core.workflow.graph_events import GraphNodeEventBase, NodeRunFailedEvent
 from core.workflow.nodes.base.node import Node
-from libs.flask_utils import preserve_flask_contexts
 
 from .ready_queue import ReadyQueue
 
+if TYPE_CHECKING:
+    pass
+
 
 @final
 class Worker(threading.Thread):
@@ -44,8 +45,7 @@ class Worker(threading.Thread):
         layers: Sequence[GraphEngineLayer],
         stop_event: threading.Event,
         worker_id: int = 0,
-        flask_app: Flask | None = None,
-        context_vars: contextvars.Context | None = None,
+        execution_context: IExecutionContext | None = None,
     ) -> None:
         """
         Initialize worker thread.
@@ -56,19 +56,17 @@ class Worker(threading.Thread):
             graph: Graph containing nodes to execute
             layers: Graph engine layers for node execution hooks
             worker_id: Unique identifier for this worker
-            flask_app: Optional Flask application for context preservation
-            context_vars: Optional context variables to preserve in worker thread
+            execution_context: Optional execution context for context preservation
         """
         super().__init__(name=f"GraphWorker-{worker_id}", daemon=True)
         self._ready_queue = ready_queue
         self._event_queue = event_queue
         self._graph = graph
         self._worker_id = worker_id
-        self._flask_app = flask_app
-        self._context_vars = context_vars
-        self._last_task_time = time.time()
+        self._execution_context = execution_context
         self._stop_event = stop_event
         self._layers = layers if layers is not None else []
+        self._last_task_time = time.time()
 
     def stop(self) -> None:
         """Worker is controlled via shared stop_event from GraphEngine.
@@ -135,11 +133,9 @@ class Worker(threading.Thread):
 
         error: Exception | None = None
 
-        if self._flask_app and self._context_vars:
-            with preserve_flask_contexts(
-                flask_app=self._flask_app,
-                context_vars=self._context_vars,
-            ):
+        # Execute the node with preserved context if execution context is provided
+        if self._execution_context is not None:
+            with self._execution_context:
                 self._invoke_node_run_start_hooks(node)
                 try:
                     node_events = node.run()

api/core/workflow/graph_engine/worker_management/worker_pool.py

@@ -8,9 +8,10 @@ DynamicScaler, and WorkerFactory into a single class.
 import logging
 import queue
 import threading
-from typing import TYPE_CHECKING, final
+from typing import final
 
 from configs import dify_config
+from core.workflow.context import IExecutionContext
 from core.workflow.graph import Graph
 from core.workflow.graph_events import GraphNodeEventBase
 
@@ -20,11 +21,6 @@ from ..worker import Worker
 
 logger = logging.getLogger(__name__)
 
-if TYPE_CHECKING:
-    from contextvars import Context
-
-    from flask import Flask
-
 
 @final
 class WorkerPool:
@@ -42,8 +38,7 @@ class WorkerPool:
         graph: Graph,
         layers: list[GraphEngineLayer],
         stop_event: threading.Event,
-        flask_app: "Flask | None" = None,
-        context_vars: "Context | None" = None,
+        execution_context: IExecutionContext | None = None,
         min_workers: int | None = None,
         max_workers: int | None = None,
         scale_up_threshold: int | None = None,
@@ -57,8 +52,7 @@ class WorkerPool:
             event_queue: Queue for worker events
             graph: The workflow graph
             layers: Graph engine layers for node execution hooks
-            flask_app: Optional Flask app for context preservation
-            context_vars: Optional context variables
+            execution_context: Optional execution context for context preservation
             min_workers: Minimum number of workers
             max_workers: Maximum number of workers
             scale_up_threshold: Queue depth to trigger scale up
@@ -67,8 +61,7 @@ class WorkerPool:
         self._ready_queue = ready_queue
         self._event_queue = event_queue
         self._graph = graph
-        self._flask_app = flask_app
-        self._context_vars = context_vars
+        self._execution_context = execution_context
         self._layers = layers
 
         # Scaling parameters with defaults
@@ -152,8 +145,7 @@ class WorkerPool:
             graph=self._graph,
             layers=self._layers,
             worker_id=worker_id,
-            flask_app=self._flask_app,
-            context_vars=self._context_vars,
+            execution_context=self._execution_context,
             stop_event=self._stop_event,
         )