mirror of
https://github.com/langgenius/dify.git
synced 2026-07-21 18:58:35 +08:00
fix(api): use resource tenant for draft variable files (#39307)
This commit is contained in:
parent
891b2dc537
commit
28d17da3b1
@ -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)
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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())
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user