fix(api): use resource tenant for draft variable files (#39307)

This commit is contained in:
林玮 (Jade Lin) 2026-07-21 10:07:47 +08:00 committed by GitHub
parent 891b2dc537
commit 28d17da3b1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 95 additions and 24 deletions

View File

@ -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)

View File

@ -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,

View File

@ -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(

View File

@ -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(

View File

@ -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,

View File

@ -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())

View File

@ -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,

View File

@ -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

View File

@ -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",

View File

@ -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",

View File

@ -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,