diff --git a/api/contexts/__init__.py b/api/contexts/__init__.py index 127b8fe76d..ae41a2c03a 100644 --- a/api/contexts/__init__.py +++ b/api/contexts/__init__.py @@ -11,10 +11,6 @@ if TYPE_CHECKING: from core.workflow.entities.variable_pool import VariablePool -tenant_id: ContextVar[str] = ContextVar("tenant_id") - -workflow_variable_pool: ContextVar["VariablePool"] = ContextVar("workflow_variable_pool") - """ To avoid race-conditions caused by gunicorn thread recycling, using RecyclableContextVar to replace with """ diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 77e315dd2d..50a843c7e8 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -158,7 +158,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): trace_manager=trace_manager, workflow_run_id=workflow_run_id, ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) @@ -240,7 +239,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): node_id=node_id, inputs=args["inputs"] ), ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) @@ -316,7 +314,6 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): extras={"auto_generate_conversation_name": False}, single_loop_run=AdvancedChatAppGenerateEntity.SingleLoopRunEntity(node_id=node_id, inputs=args["inputs"]), ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index 8b4ebede96..5442927413 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -135,7 +135,6 @@ class WorkflowAppGenerator(BaseAppGenerator): workflow_run_id=workflow_run_id, ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) @@ -282,7 +281,6 @@ class WorkflowAppGenerator(BaseAppGenerator): ), workflow_run_id=str(uuid.uuid4()), ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) @@ -359,7 +357,6 @@ class WorkflowAppGenerator(BaseAppGenerator): single_loop_run=WorkflowAppGenerateEntity.SingleLoopRunEntity(node_id=node_id, inputs=args["inputs"]), workflow_run_id=str(uuid.uuid4()), ) - contexts.tenant_id.set(application_generate_entity.app_config.tenant_id) contexts.plugin_tool_providers.set({}) contexts.plugin_tool_providers_lock.set(threading.Lock()) diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 6a5721c021..06f42494ec 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -5,7 +5,6 @@ from flask import Response, request from flask_login import user_loaded_from_request, user_logged_in from werkzeug.exceptions import NotFound, Unauthorized -import contexts from configs import dify_config from dify_app import DifyApp from extensions.ext_database import db @@ -82,8 +81,8 @@ def on_user_logged_in(_sender, user): Note: AccountService.load_logged_in_account will populate user.current_tenant_id through the load_user method, which calls account.set_tenant_id(). """ - if user and isinstance(user, Account) and user.current_tenant_id: - contexts.tenant_id.set(user.current_tenant_id) + # tenant_id context variable removed - using current_user.current_tenant_id directly + pass @login_manager.unauthorized_handler diff --git a/api/models/workflow.py b/api/models/workflow.py index b0cb8dccd9..ae341dd1b5 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -6,6 +6,8 @@ from enum import Enum, StrEnum from typing import TYPE_CHECKING, Any, Optional, Union from uuid import uuid4 +from flask_login import current_user + from core.variables import utils as variable_utils from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID from factories.variable_factory import build_segment @@ -17,7 +19,6 @@ import sqlalchemy as sa from sqlalchemy import UniqueConstraint, func from sqlalchemy.orm import Mapped, mapped_column -import contexts from constants import DEFAULT_FILE_NUMBER_LIMITS, HIDDEN_VALUE from core.helper import encrypter from core.variables import SecretVariable, Segment, SegmentType, Variable @@ -274,7 +275,16 @@ class Workflow(Base): if self._environment_variables is None: self._environment_variables = "{}" - tenant_id = contexts.tenant_id.get() + # Get tenant_id from current_user (Account or EndUser) + if isinstance(current_user, Account): + # Account user + tenant_id = current_user.current_tenant_id + else: + # EndUser + tenant_id = current_user.tenant_id + + if not tenant_id: + return [] environment_variables_dict: dict[str, Any] = json.loads(self._environment_variables) results = [ @@ -297,7 +307,17 @@ class Workflow(Base): self._environment_variables = "{}" return - tenant_id = contexts.tenant_id.get() + # Get tenant_id from current_user (Account or EndUser) + if isinstance(current_user, Account): + # Account user + tenant_id = current_user.current_tenant_id + else: + # EndUser + tenant_id = current_user.tenant_id + + if not tenant_id: + self._environment_variables = "{}" + return value = list(value) if any(var for var in value if not var.id): diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py index 70ce224eb6..6815f880a5 100644 --- a/api/tests/unit_tests/models/test_workflow.py +++ b/api/tests/unit_tests/models/test_workflow.py @@ -2,14 +2,13 @@ import json from unittest import mock from uuid import uuid4 -import contexts from constants import HIDDEN_VALUE from core.variables import FloatVariable, IntegerVariable, SecretVariable, StringVariable from models.workflow import Workflow, WorkflowNodeExecution def test_environment_variables(): - contexts.tenant_id.set("tenant_id") + # tenant_id context variable removed - using current_user.current_tenant_id directly # Create a Workflow instance workflow = Workflow( @@ -51,7 +50,7 @@ def test_environment_variables(): def test_update_environment_variables(): - contexts.tenant_id.set("tenant_id") + # tenant_id context variable removed - using current_user.current_tenant_id directly # Create a Workflow instance workflow = Workflow( @@ -104,7 +103,7 @@ def test_update_environment_variables(): def test_to_dict(): - contexts.tenant_id.set("tenant_id") + # tenant_id context variable removed - using current_user.current_tenant_id directly # Create a Workflow instance workflow = Workflow(