feat: add a flask_context_manager.

Replaces manual context management and user restoration logic with a
unified `flask_context_manager` for better code maintainability and
clarity. Removes redundant import of `has_request_context` from
multiple modules. Updates `.gitignore` to exclude `.roo/` directory,
presumably related to AI Assistant outputs.

Signed-off-by: -LAN- <laipz8200@outlook.com>

.gitignore

@@ -210,3 +210,6 @@ mise.toml
 
 # Next.js build output
 .next/
+
+# AI Assistant
+.roo/

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

@@ -5,7 +5,7 @@ import uuid
 from collections.abc import Generator, Mapping
 from typing import Any, Literal, Optional, Union, overload
 
-from flask import Flask, copy_current_request_context, current_app, has_request_context
+from flask import Flask, copy_current_request_context, current_app
 from pydantic import ValidationError
 from sqlalchemy.orm import sessionmaker
 
@@ -22,6 +22,7 @@ from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom
 from core.app.entities.task_entities import ChatbotAppBlockingResponse, ChatbotAppStreamResponse
+from core.helper import flask_context_manager
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.prompt.utils.get_thread_messages_length import get_thread_messages_length
@@ -449,24 +450,9 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
         :param message_id: message ID
         :return:
         """
-        for var, val in context.items():
-            var.set(val)
 
-        # FIXME(-LAN-): Save current user before entering new app context
+        with flask_context_manager(flask_app, context_vars=context):
-        from flask import g
-
-        saved_user = None
-        if has_request_context() and hasattr(g, "_login_user"):
-            saved_user = g._login_user
-
-        with flask_app.app_context():
             try:
-                # Restore user in new app context
-                if saved_user is not None:
-                    from flask import g
-
-                    g._login_user = saved_user
-
                 # get conversation and message
                 conversation = self._get_conversation(conversation_id)
                 message = self._get_message(message_id)

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

@@ -5,7 +5,7 @@ import uuid
 from collections.abc import Generator, Mapping
 from typing import Any, Literal, Union, overload
 
-from flask import Flask, copy_current_request_context, current_app, has_request_context
+from flask import Flask, copy_current_request_context, current_app
 from pydantic import ValidationError
 
 from configs import dify_config
@@ -19,6 +19,7 @@ from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskSt
 from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager
 from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, InvokeFrom
+from core.helper import flask_context_manager
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.ops.ops_trace_manager import TraceQueueManager
 from extensions.ext_database import db
@@ -229,24 +230,9 @@ class AgentChatAppGenerator(MessageBasedAppGenerator):
         :param message_id: message ID
         :return:
         """
-        for var, val in context.items():
-            var.set(val)
 
-        # FIXME(-LAN-): Save current user before entering new app context
+        with flask_context_manager(flask_app, context_vars=context):
-        from flask import g
-
-        saved_user = None
-        if has_request_context() and hasattr(g, "_login_user"):
-            saved_user = g._login_user
-
-        with flask_app.app_context():
             try:
-                # Restore user in new app context
-                if saved_user is not None:
-                    from flask import g
-
-                    g._login_user = saved_user
-
                 # get conversation and message
                 conversation = self._get_conversation(conversation_id)
                 message = self._get_message(message_id)

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

@@ -5,7 +5,7 @@ import uuid
 from collections.abc import Generator, Mapping, Sequence
 from typing import Any, Literal, Optional, Union, overload
 
-from flask import Flask, copy_current_request_context, current_app, has_request_context
+from flask import Flask, copy_current_request_context, current_app
 from pydantic import ValidationError
 from sqlalchemy.orm import sessionmaker
 
@@ -21,6 +21,7 @@ from core.app.apps.workflow.generate_response_converter import WorkflowAppGenera
 from core.app.apps.workflow.generate_task_pipeline import WorkflowAppGenerateTaskPipeline
 from core.app.entities.app_invoke_entities import InvokeFrom, WorkflowAppGenerateEntity
 from core.app.entities.task_entities import WorkflowAppBlockingResponse, WorkflowAppStreamResponse
+from core.helper import flask_context_manager
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository
@@ -408,24 +409,9 @@ class WorkflowAppGenerator(BaseAppGenerator):
         :param workflow_thread_pool_id: workflow thread pool id
         :return:
         """
-        for var, val in context.items():
-            var.set(val)
 
-        # FIXME(-LAN-): Save current user before entering new app context
+        with flask_context_manager(flask_app, context_vars=context):
-        from flask import g
-
-        saved_user = None
-        if has_request_context() and hasattr(g, "_login_user"):
-            saved_user = g._login_user
-
-        with flask_app.app_context():
             try:
-                # Restore user in new app context
-                if saved_user is not None:
-                    from flask import g
-
-                    g._login_user = saved_user
-
                 # workflow app
                 runner = WorkflowAppRunner(
                     application_generate_entity=application_generate_entity,

api/core/helper/__init__.py

@@ -0,0 +1,3 @@
+from core.helper.flask_context import flask_context_manager
+
+__all__ = ["flask_context_manager"]

api/core/helper/flask_context.py

@@ -0,0 +1,63 @@
+import contextvars
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from typing import Any, Optional, TypeVar, Union
+
+from flask import Flask, g, has_request_context
+
+T = TypeVar("T")
+
+
+@contextmanager
+def flask_context_manager(
+    flask_app: Flask,
+    context_vars: Optional[
+        Union[Mapping[contextvars.ContextVar[T], Any], dict[contextvars.ContextVar[Any], Any]]
+    ] = None,
+) -> Iterator[None]:
+    """
+    A context manager that handles:
+    1. flask-login's UserProxy copy
+    2. ContextVars copy
+    3. flask_app.app_context()
+
+    This context manager ensures that the Flask application context is properly set up,
+    the current user is preserved across context boundaries, and any provided context variables
+    are set within the new context.
+
+    Args:
+        flask_app: The Flask application instance
+        context_vars: Optional dictionary mapping ContextVar instances to their values
+
+    Yields:
+        None
+
+    Example:
+        ```python
+        with flask_context_manager(flask_app, context_vars={my_var: my_value}):
+            # Code that needs Flask app context and context variables
+            # Current user will be preserved if available
+        ```
+    """
+    # Set context variables if provided
+    if context_vars:
+        for var, val in context_vars.items():
+            var.set(val)
+
+    # Save current user before entering new app context
+    saved_user = None
+    if has_request_context() and hasattr(g, "_login_user"):
+        saved_user = g._login_user
+
+    # Enter Flask app context
+    with flask_app.app_context():
+        try:
+            # Restore user in new app context if it was saved
+            if saved_user is not None:
+                g._login_user = saved_user
+
+            # Yield control back to the caller
+            yield
+        finally:
+            # Any cleanup can be added here if needed
+            pass

api/core/workflow/graph_engine/graph_engine.py

@@ -9,11 +9,12 @@ from copy import copy, deepcopy
 from datetime import UTC, datetime
 from typing import Any, Optional, cast
 
-from flask import Flask, current_app, has_request_context
+from flask import Flask, current_app
 
 from configs import dify_config
 from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
+from core.helper import flask_context_manager
 from core.workflow.entities.node_entities import AgentNodeStrategyInit, NodeRunResult
 from core.workflow.entities.variable_pool import VariablePool, VariableValue
 from core.workflow.entities.workflow_node_execution import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
@@ -537,24 +538,9 @@ class GraphEngine:
         """
         Run parallel nodes
         """
-        for var, val in context.items():
-            var.set(val)
 
-        # FIXME(-LAN-): Save current user before entering new app context
+        with flask_context_manager(flask_app, context_vars=context):
-        from flask import g
-
-        saved_user = None
-        if has_request_context() and hasattr(g, "_login_user"):
-            saved_user = g._login_user
-
-        with flask_app.app_context():
             try:
-                # Restore user in new app context
-                if saved_user is not None:
-                    from flask import g
-
-                    g._login_user = saved_user
-
                 q.put(
                     ParallelBranchRunStartedEvent(
                         parallel_id=parallel_id,

api/core/workflow/nodes/iteration/iteration_node.py

@@ -7,9 +7,10 @@ from datetime import UTC, datetime
 from queue import Empty, Queue
 from typing import TYPE_CHECKING, Any, Optional, cast
 
-from flask import Flask, current_app, has_request_context
+from flask import Flask, current_app
 
 from configs import dify_config
+from core.helper import flask_context_manager
 from core.variables import ArrayVariable, IntegerVariable, NoneVariable
 from core.workflow.entities.node_entities import (
     NodeRunResult,
@@ -583,23 +584,8 @@ class IterationNode(BaseNode[IterationNodeData]):
         """
         run single iteration in parallel mode
         """
-        for var, val in context.items():
-            var.set(val)
-
-        # FIXME(-LAN-): Save current user before entering new app context
-        from flask import g
-
-        saved_user = None
-        if has_request_context() and hasattr(g, "_login_user"):
-            saved_user = g._login_user
-
-        with flask_app.app_context():
-            # Restore user in new app context
-            if saved_user is not None:
-                from flask import g
-
-                g._login_user = saved_user
 
+        with flask_context_manager(flask_app, context_vars=context):
             parallel_mode_run_id = uuid.uuid4().hex
             graph_engine_copy = graph_engine.create_copy()
             variable_pool_copy = graph_engine_copy.graph_runtime_state.variable_pool