From 225c5fe2a3efea0f735224316d0a179aeee35cfe Mon Sep 17 00:00:00 2001 From: -LAN- Date: Mon, 16 Jun 2025 18:09:34 +0800 Subject: [PATCH] refactor: Renames flask_context_manager to preserve_flask_contexts Signed-off-by: -LAN- --- .../app/apps/advanced_chat/app_generator.py | 4 ++-- api/core/app/apps/agent_chat/app_generator.py | 4 ++-- api/core/app/apps/workflow/app_generator.py | 4 ++-- api/core/workflow/graph_engine/graph_engine.py | 4 ++-- .../workflow/nodes/iteration/iteration_node.py | 4 ++-- api/libs/flask_utils.py | 4 ++-- api/tests/unit_tests/libs/test_flask_utils.py | 18 +++++++++--------- 7 files changed, 21 insertions(+), 21 deletions(-) diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 805de72e46..9e6adc4b08 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -31,7 +31,7 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository from extensions.ext_database import db from factories import file_factory -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts from models import Account, App, Conversation, EndUser, Message, Workflow, WorkflowNodeExecutionTriggeredFrom from models.enums import WorkflowRunTriggeredFrom from services.conversation_service import ConversationService @@ -448,7 +448,7 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): :return: """ - with flask_context_manager(flask_app, context_vars=context): + with preserve_flask_contexts(flask_app, context_vars=context): try: # get conversation and message conversation = self._get_conversation(conversation_id) diff --git a/api/core/app/apps/agent_chat/app_generator.py b/api/core/app/apps/agent_chat/app_generator.py index 5279155504..a448bf8a94 100644 --- a/api/core/app/apps/agent_chat/app_generator.py +++ b/api/core/app/apps/agent_chat/app_generator.py @@ -23,7 +23,7 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError from core.ops.ops_trace_manager import TraceQueueManager from extensions.ext_database import db from factories import file_factory -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts from models import Account, App, EndUser from services.conversation_service import ConversationService from services.errors.message import MessageNotExistsError @@ -228,7 +228,7 @@ class AgentChatAppGenerator(MessageBasedAppGenerator): :return: """ - with flask_context_manager(flask_app, context_vars=context): + with preserve_flask_contexts(flask_app, context_vars=context): try: # get conversation and message conversation = self._get_conversation(conversation_id) diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index 8ca7510a6b..7f4770fc97 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -29,7 +29,7 @@ from core.workflow.repositories.workflow_execution_repository import WorkflowExe from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository from extensions.ext_database import db from factories import file_factory -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts from models import Account, App, EndUser, Workflow, WorkflowNodeExecutionTriggeredFrom from models.enums import WorkflowRunTriggeredFrom @@ -407,7 +407,7 @@ class WorkflowAppGenerator(BaseAppGenerator): :return: """ - with flask_context_manager(flask_app, context_vars=context): + with preserve_flask_contexts(flask_app, context_vars=context): try: # workflow app runner = WorkflowAppRunner( diff --git a/api/core/workflow/graph_engine/graph_engine.py b/api/core/workflow/graph_engine/graph_engine.py index 5cb123a961..875cee17e6 100644 --- a/api/core/workflow/graph_engine/graph_engine.py +++ b/api/core/workflow/graph_engine/graph_engine.py @@ -53,7 +53,7 @@ from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts from models.enums import UserFrom from models.workflow import WorkflowType @@ -539,7 +539,7 @@ class GraphEngine: Run parallel nodes """ - with flask_context_manager(flask_app, context_vars=context): + with preserve_flask_contexts(flask_app, context_vars=context): try: q.put( ParallelBranchRunStartedEvent( diff --git a/api/core/workflow/nodes/iteration/iteration_node.py b/api/core/workflow/nodes/iteration/iteration_node.py index becb56eabc..42b6795fb0 100644 --- a/api/core/workflow/nodes/iteration/iteration_node.py +++ b/api/core/workflow/nodes/iteration/iteration_node.py @@ -37,7 +37,7 @@ from core.workflow.nodes.base import BaseNode from core.workflow.nodes.enums import NodeType from core.workflow.nodes.event import NodeEvent, RunCompletedEvent from core.workflow.nodes.iteration.entities import ErrorHandleMode, IterationNodeData -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts from .exc import ( InvalidIteratorValueError, @@ -585,7 +585,7 @@ class IterationNode(BaseNode[IterationNodeData]): run single iteration in parallel mode """ - with flask_context_manager(flask_app, context_vars=context): + with preserve_flask_contexts(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 diff --git a/api/libs/flask_utils.py b/api/libs/flask_utils.py index 6460b8e02f..4ea2779584 100644 --- a/api/libs/flask_utils.py +++ b/api/libs/flask_utils.py @@ -9,7 +9,7 @@ T = TypeVar("T") @contextmanager -def flask_context_manager( +def preserve_flask_contexts( flask_app: Flask, context_vars: contextvars.Context, ) -> Iterator[None]: @@ -36,7 +36,7 @@ def flask_context_manager( Example: ```python - with flask_context_manager(flask_app, context_vars=context_vars): + with preserve_flask_contexts(flask_app, context_vars=context_vars): # Code that needs Flask app context and context variables # Current user will be preserved if available ``` diff --git a/api/tests/unit_tests/libs/test_flask_utils.py b/api/tests/unit_tests/libs/test_flask_utils.py index b2f8fa7e56..fb46ba50f3 100644 --- a/api/tests/unit_tests/libs/test_flask_utils.py +++ b/api/tests/unit_tests/libs/test_flask_utils.py @@ -6,7 +6,7 @@ import pytest from flask import Flask from flask_login import LoginManager, UserMixin, current_user, login_user -from libs.flask_utils import flask_context_manager +from libs.flask_utils import preserve_flask_contexts class User(UserMixin): @@ -45,9 +45,9 @@ def test_user() -> User: def test_current_user_not_accessible_across_threads(login_app: Flask, test_user: User): """ - Test that current_user is not accessible in a different thread without flask_context_manager. + Test that current_user is not accessible in a different thread without preserve_flask_contexts. - This test demonstrates that without the flask_context_manager, we cannot access + This test demonstrates that without the preserve_flask_contexts, we cannot access current_user in a different thread, even with app_context. """ # Log in the user in the main thread @@ -65,7 +65,7 @@ def test_current_user_not_accessible_across_threads(login_app: Flask, test_user: # Try to access current_user in a different thread with app_context with login_app.app_context(): # This should fail because current_user is not accessible across threads - # without flask_context_manager + # without preserve_flask_contexts result["user_accessible"] = current_user.is_authenticated except Exception as e: result["error"] = str(e) # type: ignore @@ -79,11 +79,11 @@ def test_current_user_not_accessible_across_threads(login_app: Flask, test_user: assert result["error"] is not None or (result["user_accessible"] is not None and not result["user_accessible"]) -def test_current_user_accessible_with_flask_context_manager(login_app: Flask, test_user: User): +def test_current_user_accessible_with_preserve_flask_contexts(login_app: Flask, test_user: User): """ - Test that current_user is accessible in a different thread with flask_context_manager. + Test that current_user is accessible in a different thread with preserve_flask_contexts. - This test demonstrates that with the flask_context_manager, we can access + This test demonstrates that with the preserve_flask_contexts, we can access current_user in a different thread. """ # Log in the user in the main thread @@ -101,8 +101,8 @@ def test_current_user_accessible_with_flask_context_manager(login_app: Flask, te # Define a function to run in a separate thread def check_user_in_thread_with_manager(): try: - # Use flask_context_manager to access current_user in a different thread - with flask_context_manager(login_app, context_vars): + # Use preserve_flask_contexts to access current_user in a different thread + with preserve_flask_contexts(login_app, context_vars): from flask_login import current_user if current_user: