From 28d17da3b1bf35444caa820bfd128dc42e98b469 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=97=E7=8E=AE=20=28Jade=20Lin=29?= Date: Tue, 21 Jul 2026 10:07:47 +0800 Subject: [PATCH] fix(api): use resource tenant for draft variable files (#39307) --- .../app/apps/advanced_chat/app_generator.py | 6 ++- api/core/app/apps/base_app_generator.py | 11 +++- .../app/apps/pipeline/pipeline_generator.py | 1 + api/core/app/apps/workflow/app_generator.py | 6 ++- api/services/rag_pipeline/rag_pipeline.py | 2 + .../workflow_draft_variable_service.py | 21 ++++---- api/services/workflow_service.py | 2 + .../test_workflow_draft_variable_service.py | 1 + .../core/app/apps/test_base_app_generator.py | 6 ++- .../app/apps/test_workflow_app_generator.py | 11 ++-- .../test_workflow_draft_variable_service.py | 52 ++++++++++++++++--- 11 files changed, 95 insertions(+), 24 deletions(-) diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 42233b22477..3c677804028 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -625,7 +625,11 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): message=message_snapshot, user=user, stream=stream, - draft_var_saver_factory=self._get_draft_var_saver_factory(invoke_from, account=user), + draft_var_saver_factory=self._get_draft_var_saver_factory( + invoke_from, + account=user, + tenant_id=application_generate_entity.app_config.tenant_id, + ), ) return AdvancedChatAppGenerateResponseConverter.convert(response=response, invoke_from=invoke_from) diff --git a/api/core/app/apps/base_app_generator.py b/api/core/app/apps/base_app_generator.py index d7ef5165f03..2762f99301d 100644 --- a/api/core/app/apps/base_app_generator.py +++ b/api/core/app/apps/base_app_generator.py @@ -32,6 +32,7 @@ class _DebuggerDraftVariableSaver: self, *, account: Account, + tenant_id: str, app_id: str, node_id: str, node_type: NodeType, @@ -39,6 +40,7 @@ class _DebuggerDraftVariableSaver: enclosing_node_id: str | None = None, ) -> None: self._account = account + self._tenant_id = tenant_id self._app_id = app_id self._node_id = node_id self._node_type = node_type @@ -49,6 +51,7 @@ class _DebuggerDraftVariableSaver: with Session(db.engine) as session, session.begin(): DraftVariableSaverImpl( session=session, + tenant_id=self._tenant_id, app_id=self._app_id, node_id=self._node_id, node_type=self._node_type, @@ -287,7 +290,12 @@ class BaseAppGenerator: @final @staticmethod - def _get_draft_var_saver_factory(invoke_from: InvokeFrom, account: Account | EndUser) -> DraftVariableSaverFactory: + def _get_draft_var_saver_factory( + invoke_from: InvokeFrom, + account: Account | EndUser, + *, + tenant_id: str, + ) -> DraftVariableSaverFactory: if invoke_from == InvokeFrom.DEBUGGER: assert isinstance(account, Account) @@ -300,6 +308,7 @@ class BaseAppGenerator: ) -> DraftVariableSaver: return _DebuggerDraftVariableSaver( account=account, + tenant_id=tenant_id, app_id=app_id, node_id=node_id, node_type=node_type, diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index dca2eb9055f..3eb93e7c08a 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/api/core/app/apps/pipeline/pipeline_generator.py @@ -349,6 +349,7 @@ class PipelineGenerator(BaseAppGenerator): draft_var_saver_factory = self._get_draft_var_saver_factory( invoke_from, user, + tenant_id=pipeline.tenant_id, ) # return response or stream generator response = self._handle_response( diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index 3dc2cf99e04..fb5393d7730 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -399,7 +399,11 @@ class WorkflowAppGenerator(BaseAppGenerator): worker_thread.start() - draft_var_saver_factory = self._get_draft_var_saver_factory(invoke_from, user) + draft_var_saver_factory = self._get_draft_var_saver_factory( + invoke_from, + user, + tenant_id=app_model.tenant_id, + ) # return response or stream generator response = self._handle_response( diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 6d7f3a01af4..993f6d79492 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -601,6 +601,7 @@ class RagPipelineService: with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, + tenant_id=pipeline.tenant_id, app_id=pipeline.id, node_id=workflow_node_execution.node_id, node_type=workflow_node_execution.node_type, @@ -1391,6 +1392,7 @@ class RagPipelineService: with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, + tenant_id=pipeline.tenant_id, app_id=pipeline.id, node_id=workflow_node_execution_db_model.node_id, node_type=workflow_node_execution_db_model.node_type, diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index 9c30f32a7ad..e0a0c18d5c6 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -823,6 +823,8 @@ _FILENAME_TRANS_TABLE = _make_filename_trans_table() class DraftVariableSaver: + """Persist draft outputs under the tenant that owns the app or pipeline.""" + # _DUMMY_OUTPUT_IDENTITY is a placeholder output for workflow nodes. # Its sole possible value is `None`. # @@ -842,6 +844,10 @@ class DraftVariableSaver: # Database session used for persisting draft variables. _session: Session + # Resource owner tenant. An account's current tenant may be unset or point elsewhere + # when draft variables are persisted by an asynchronous workflow execution. + _tenant_id: str + # The application ID associated with the draft variables. # This should match the `Workflow.app_id` of the workflow to which the current node belongs. _app_id: str @@ -867,6 +873,7 @@ class DraftVariableSaver: def __init__( self, session: Session, + tenant_id: str, app_id: str, node_id: str, node_type: NodeType, @@ -878,6 +885,7 @@ class DraftVariableSaver: # WorkflowNodeExecutionModel/WorkflowNodeExecution, not their `node_execution_id` # field. These are distinct database fields with different purposes. self._session = session + self._tenant_id = tenant_id self._app_id = app_id self._node_id = node_id self._node_type = node_type @@ -885,12 +893,6 @@ class DraftVariableSaver: self._user = user self._enclosing_node_id = enclosing_node_id - def _resolve_app_tenant_id(self) -> str: - tenant_id = self._session.scalar(select(App.tenant_id).where(App.id == self._app_id)) - if not tenant_id: - raise ValueError(f"Unable to resolve tenant_id for app {self._app_id}") - return tenant_id - def _create_dummy_output_variable(self): return WorkflowDraftVariable.new_node_variable( app_id=self._app_id, @@ -949,11 +951,10 @@ class DraftVariableSaver: if name == SystemVariableKey.FILES: # Here we know the type of variable must be `array[file]`, we # just rebuild files from the serialized payload. - tenant_id = self._resolve_app_tenant_id() files = [ build_file_from_stored_mapping( file_mapping=v, - tenant_id=tenant_id, + tenant_id=self._tenant_id, ) for v in value ] @@ -1096,8 +1097,8 @@ class DraftVariableSaver: content=original_content_serialized.encode(), mimetype=content_type, user=self._user, + tenant_id=self._tenant_id, ) - assert self._user.current_tenant_id # Create WorkflowDraftVariableFile record variable_file = WorkflowDraftVariableFile( upload_file_id=upload_file.id, @@ -1105,7 +1106,7 @@ class DraftVariableSaver: length=original_length, value_type=value_seg.value_type, app_id=self._app_id, - tenant_id=self._user.current_tenant_id, + tenant_id=self._tenant_id, user_id=self._user.id, ) variable_file.id = str(uuidv7()) diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index b29affe696a..b426f3363c0 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -1060,6 +1060,7 @@ class WorkflowService: with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, + tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=workflow_node_execution.node_id, node_type=workflow_node_execution.node_type, @@ -1210,6 +1211,7 @@ class WorkflowService: with sessionmaker(bind=db.engine).begin() as session: draft_var_saver = DraftVariableSaver( session=session, + tenant_id=app_model.tenant_id, app_id=app_model.id, node_id=node_id, node_type=BuiltinNodeTypes.HUMAN_INPUT, diff --git a/api/tests/integration_tests/services/test_workflow_draft_variable_service.py b/api/tests/integration_tests/services/test_workflow_draft_variable_service.py index 737f3ff5806..37bcdce6334 100644 --- a/api/tests/integration_tests/services/test_workflow_draft_variable_service.py +++ b/api/tests/integration_tests/services/test_workflow_draft_variable_service.py @@ -311,6 +311,7 @@ class TestDraftVariableLoader(unittest.TestCase): # Use DraftVariableSaver to create offloaded variable (this mimics production) saver = DraftVariableSaver( session=session, + tenant_id=self._test_tenant_id, app_id=self._test_app_id, node_id="test_offload_node", node_type=BuiltinNodeTypes.LLM, # Use a real node type diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py b/api/tests/unit_tests/core/app/apps/test_base_app_generator.py index f2a17006645..8e7468bb0b8 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_generator.py @@ -485,7 +485,11 @@ class TestBaseAppGeneratorExtras: account.id = "account-id" account.tenant_id = "tenant-id" - factory = base_app_generator._get_draft_var_saver_factory(InvokeFrom.DEBUGGER, account) + factory = base_app_generator._get_draft_var_saver_factory( + InvokeFrom.DEBUGGER, + account, + tenant_id="tenant-id", + ) saver = factory( app_id="app-id", node_id="node-id", diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py index 50c5a6811db..8f5cb2b8115 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_generator.py @@ -185,7 +185,11 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture): return_value="converted", ) mocker.patch.object(WorkflowAppGenerator, "_handle_response", return_value="response") - mocker.patch.object(WorkflowAppGenerator, "_get_draft_var_saver_factory", return_value=MagicMock()) + draft_saver_factory = mocker.patch.object( + WorkflowAppGenerator, + "_get_draft_var_saver_factory", + return_value=MagicMock(), + ) pause_layer = MagicMock(name="pause-layer") mocker.patch( @@ -209,7 +213,7 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture): mocker.patch("core.app.apps.workflow.app_generator.threading.Thread", DummyThread) - app_model = SimpleNamespace(mode="workflow") + app_model = SimpleNamespace(mode="workflow", tenant_id="tenant") app_config = SimpleNamespace(app_id="app", tenant_id="tenant", workflow_id="wf") application_generate_entity = SimpleNamespace( task_id="task", @@ -240,6 +244,7 @@ def test_generate_appends_pause_layer_and_forwards_state(mocker: MockerFixture): assert result == "converted" assert worker_kwargs["kwargs"]["graph_engine_layers"] == ("base-layer", pause_layer) assert worker_kwargs["kwargs"]["graph_runtime_state"] is graph_runtime_state + assert draft_saver_factory.call_args.kwargs["tenant_id"] == app_model.tenant_id def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture): @@ -301,7 +306,7 @@ def test_resume_path_runs_worker_with_runtime_state(mocker: MockerFixture): pause_config = SimpleNamespace(session_factory=MagicMock(), state_owner_user_id="owner") - app_model = SimpleNamespace(mode="workflow") + app_model = SimpleNamespace(mode="workflow", tenant_id="tenant") app_config = SimpleNamespace(app_id="app", tenant_id="tenant", workflow_id="workflow") application_generate_entity = SimpleNamespace( task_id="task", diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py index 5fb83a1cf86..c041a6203e9 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py @@ -1,6 +1,7 @@ import dataclasses import secrets import uuid +from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch import pytest @@ -59,6 +60,7 @@ class TestDraftVariableSaver: test_app_id = self._get_test_app_id() saver = DraftVariableSaver( session=mock_session, + tenant_id="test-tenant-id", app_id=test_app_id, node_id="test_node_id", node_type=BuiltinNodeTypes.START, @@ -121,6 +123,7 @@ class TestDraftVariableSaver: test_app_id = self._get_test_app_id() saver = DraftVariableSaver( session=mock_session, + tenant_id="test-tenant-id", app_id=test_app_id, node_id=_NODE_ID, node_type=BuiltinNodeTypes.START, @@ -139,6 +142,7 @@ class TestDraftVariableSaver: mock_user.id = str(uuid.uuid4()) saver = DraftVariableSaver( session=mock_session, + tenant_id="tenant-1", app_id=self._get_test_app_id(), node_id="start", node_type=BuiltinNodeTypes.START, @@ -161,13 +165,10 @@ class TestDraftVariableSaver: "tenant_id": "legacy-tenant", } - with ( - patch.object(saver, "_resolve_app_tenant_id", return_value="tenant-1"), - patch( - "services.workflow_draft_variable_service.build_file_from_stored_mapping", - return_value=rebuilt_file, - ) as rebuild_file, - ): + with patch( + "services.workflow_draft_variable_service.build_file_from_stored_mapping", + return_value=rebuilt_file, + ) as rebuild_file: draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]}) sys_var = draft_vars[0] @@ -194,6 +195,7 @@ class TestDraftVariableSaver: return DraftVariableSaver( session=mock_session, + tenant_id="test-tenant-id", app_id="test-app-id", node_id="test-node-id", node_type=BuiltinNodeTypes.LLM, @@ -235,6 +237,40 @@ class TestDraftVariableSaver: # Should not have large variable metadata assert draft_var.file_id == mock_draft_var_file.id + def test_try_offload_large_variable_uses_resource_tenant(self, mock_session): + mock_user = MagicMock(spec=Account) + mock_user.id = "test-user-id" + mock_user.current_tenant_id = "" + saver = DraftVariableSaver( + session=mock_session, + tenant_id="app-tenant-id", + app_id="test-app-id", + node_id="test-node-id", + node_type=BuiltinNodeTypes.LLM, + node_execution_id="test-execution-id", + user=mock_user, + ) + upload_file = SimpleNamespace(id="upload-file-id") + truncation_result = SimpleNamespace(result=StringSegment(value="..."), truncated=True) + + with ( + patch( + "services.workflow_draft_variable_service.VariableTruncator.truncate", return_value=truncation_result + ), + patch("services.workflow_draft_variable_service.FileService") as file_service_class, + patch("services.workflow_draft_variable_service.sessionmaker") as sessionmaker_mock, + ): + file_service_class.return_value.upload_file.return_value = upload_file + result = saver._try_offload_large_variable("large_var", StringSegment(value="large value")) + + assert result is not None + _, variable_file = result + assert file_service_class.return_value.upload_file.call_args.kwargs["tenant_id"] == "app-tenant-id" + assert variable_file.tenant_id == "app-tenant-id" + sessionmaker_mock.return_value.begin.return_value.__enter__.return_value.add.assert_called_once_with( + variable_file + ) + @patch("services.workflow_draft_variable_service._batch_upsert_draft_variable", autospec=True) def test_save_method_integration(self, mock_batch_upsert, draft_saver): """Test complete save workflow.""" @@ -257,6 +293,7 @@ class TestDraftVariableSaver: saver = DraftVariableSaver( session=mock_session, + tenant_id="test-tenant-id", app_id="test-app-id", node_id="start-node-id", node_type=BuiltinNodeTypes.START, @@ -292,6 +329,7 @@ class TestDraftVariableSaver: saver = DraftVariableSaver( session=mock_session, + tenant_id="test-tenant-id", app_id="test-app-id", node_id="start-node-id", node_type=BuiltinNodeTypes.START,