diff --git a/.gitignore b/.gitignore index 74a9ef63ef..486a5e42b0 100644 --- a/.gitignore +++ b/.gitignore @@ -210,3 +210,6 @@ mise.toml # Next.js build output .next/ + +# AI Assistant +.roo/ \ No newline at end of file diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 8c85f91d7e..ef43fdd7d2 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/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 - 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(): + with flask_context_manager(flask_app, context_vars=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) diff --git a/api/core/app/apps/agent_chat/app_generator.py b/api/core/app/apps/agent_chat/app_generator.py index 158196f24d..f3a85d3782 100644 --- a/api/core/app/apps/agent_chat/app_generator.py +++ b/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 - 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(): + with flask_context_manager(flask_app, context_vars=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) diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index f4aec3479b..c36f4b86cc 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/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 - 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(): + with flask_context_manager(flask_app, context_vars=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, diff --git a/api/core/helper/__init__.py b/api/core/helper/__init__.py index e69de29bb2..a769d5d1ed 100644 --- a/api/core/helper/__init__.py +++ b/api/core/helper/__init__.py @@ -0,0 +1,3 @@ +from core.helper.flask_context import flask_context_manager + +__all__ = ["flask_context_manager"] diff --git a/api/core/helper/flask_context.py b/api/core/helper/flask_context.py new file mode 100644 index 0000000000..cb4b94d33c --- /dev/null +++ b/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 diff --git a/api/core/workflow/graph_engine/graph_engine.py b/api/core/workflow/graph_engine/graph_engine.py index 363b2ee920..255f14a89c 100644 --- a/api/core/workflow/graph_engine/graph_engine.py +++ b/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 - 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(): + with flask_context_manager(flask_app, context_vars=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, diff --git a/api/core/workflow/nodes/iteration/iteration_node.py b/api/core/workflow/nodes/iteration/iteration_node.py index 2592823540..83f05204bd 100644 --- a/api/core/workflow/nodes/iteration/iteration_node.py +++ b/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