diff --git a/api/tasks/delete_conversation_task.py b/api/tasks/delete_conversation_task.py index 38326b1bacb..c9bf4ce9f17 100644 --- a/api/tasks/delete_conversation_task.py +++ b/api/tasks/delete_conversation_task.py @@ -3,7 +3,7 @@ import time import click from celery import shared_task -from sqlalchemy import delete, select, update +from sqlalchemy import delete, select from configs import dify_config from core.db.session_factory import session_factory @@ -50,7 +50,8 @@ def _cleanup_conversation_related_data(conversation_id: str) -> bool: The storage object is deleted before its ``ToolFile`` row so a failed attempt retains the durable ``file_key`` needed by the next retry. ToolFiles promoted - to Agent Drive are detached from the conversation and survive this lifecycle. + to Agent Drive are detached from the conversation, and their Drive references + take over lifecycle ownership. """ with session_factory.create_session() as session: @@ -59,18 +60,6 @@ def _cleanup_conversation_related_data(conversation_id: str) -> bool: logger.warning("Skipped cleanup for active conversation %s", conversation_id) return False - drive_tool_file_ids = select(AgentDriveFile.file_id).where( - AgentDriveFile.file_kind == AgentDriveFileKind.TOOL_FILE - ) - session.execute( - update(ToolFile) - .where( - ToolFile.conversation_id == conversation_id, - ToolFile.id.in_(drive_tool_file_ids), - ) - .values(conversation_id=None) - ) - tool_files = list( session.scalars( select(ToolFile) @@ -79,7 +68,25 @@ def _cleanup_conversation_related_data(conversation_id: str) -> bool: .with_for_update() ) ) + tool_file_ids = [tool_file.id for tool_file in tool_files] + drive_files = list( + session.scalars( + select(AgentDriveFile) + .where( + AgentDriveFile.file_kind == AgentDriveFileKind.TOOL_FILE, + AgentDriveFile.file_id.in_(tool_file_ids), + ) + .order_by(AgentDriveFile.id) + .with_for_update() + ) + ) + drive_tool_file_ids = {drive_file.file_id for drive_file in drive_files} + for drive_file in drive_files: + drive_file.value_owned_by_drive = True for tool_file in tool_files: + if tool_file.id in drive_tool_file_ids: + tool_file.conversation_id = None + continue _delete_storage_object(tool_file.file_key) session.delete(tool_file) diff --git a/api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py b/api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py index 015b618aea4..19d384b0491 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py @@ -1,5 +1,7 @@ +from threading import Event, Thread from unittest.mock import patch +import sqlalchemy as sa from sqlalchemy import select from sqlalchemy.orm import Session @@ -55,7 +57,7 @@ def test_cleanup_deletes_owned_storage_and_preserves_drive_file( key="drive.txt", file_kind=AgentDriveFileKind.TOOL_FILE, file_id=drive_file.id, - value_owned_by_drive=True, + value_owned_by_drive=False, is_skill=False, ) db_session_with_containers.add(drive_entry) @@ -73,6 +75,105 @@ def test_cleanup_deletes_owned_storage_and_preserves_drive_file( preserved = db_session_with_containers.get(ToolFile, drive_file_id) assert preserved is not None assert preserved.conversation_id is None - assert db_session_with_containers.scalar( + preserved_drive_entry = db_session_with_containers.scalar( select(AgentDriveFile).where(AgentDriveFile.file_id == drive_file_id) - ) is not None + ) + assert preserved_drive_entry is not None + assert preserved_drive_entry.value_owned_by_drive is True + + +def test_cleanup_preserves_drive_file_committed_while_waiting_for_tool_file_lock( + db_session_with_containers: Session, +) -> None: + conversation = Conversation( + id=CONVERSATION_ID, + app_id=APP_ID, + mode=AppMode.CHAT, + name="Deleted conversation", + inputs={}, + status=ConversationStatus.NORMAL, + from_source=ConversationFromSource.CONSOLE, + from_account_id=ACCOUNT_ID, + is_deleted=True, + ) + drive_file = ToolFile( + user_id=ACCOUNT_ID, + tenant_id=TENANT_ID, + conversation_id=CONVERSATION_ID, + file_key=f"tools/{TENANT_ID}/concurrent-drive.txt", + mimetype="text/plain", + name="concurrent-drive.txt", + size=5, + ) + db_session_with_containers.add_all([conversation, drive_file]) + db_session_with_containers.commit() + drive_file_id = drive_file.id + + engine = db_session_with_containers.get_bind() + drive_session = Session(engine) + locked_file = drive_session.scalar(select(ToolFile).where(ToolFile.id == drive_file_id).with_for_update()) + assert locked_file is not None + drive_session.add( + AgentDriveFile( + tenant_id=TENANT_ID, + agent_id=AGENT_ID, + key="concurrent-drive.txt", + file_kind=AgentDriveFileKind.TOOL_FILE, + file_id=drive_file_id, + value_owned_by_drive=False, + is_skill=False, + ) + ) + drive_session.flush() + + cleanup_result: list[bool] = [] + cleanup_errors: list[BaseException] = [] + + def run_cleanup() -> None: + try: + cleanup_result.append(_cleanup_conversation_related_data(CONVERSATION_ID)) + except BaseException as error: + cleanup_errors.append(error) + + tool_file_lock_started = Event() + + def signal_tool_file_lock( + _connection, + _cursor, + statement: str, + _parameters, + _context, + _executemany, + ) -> None: + normalized_statement = statement.lower() + if "from tool_files" in normalized_statement and "for update" in normalized_statement: + tool_file_lock_started.set() + + sa.event.listen(engine, "before_cursor_execute", signal_tool_file_lock) + cleanup_thread = Thread(target=run_cleanup) + try: + with patch("tasks.delete_conversation_task.storage") as storage_mock: + cleanup_thread.start() + assert tool_file_lock_started.wait(timeout=5) + drive_session.commit() + cleanup_thread.join(timeout=5) + finally: + sa.event.remove(engine, "before_cursor_execute", signal_tool_file_lock) + drive_session.rollback() + drive_session.close() + cleanup_thread.join(timeout=5) + + assert not cleanup_thread.is_alive() + assert cleanup_errors == [] + assert cleanup_result == [True] + storage_mock.delete.assert_not_called() + + db_session_with_containers.expire_all() + preserved = db_session_with_containers.get(ToolFile, drive_file_id) + assert preserved is not None + assert preserved.conversation_id is None + preserved_drive_entry = db_session_with_containers.scalar( + select(AgentDriveFile).where(AgentDriveFile.file_id == drive_file_id) + ) + assert preserved_drive_entry is not None + assert preserved_drive_entry.value_owned_by_drive is True diff --git a/api/tests/unit_tests/tasks/test_delete_conversation_task.py b/api/tests/unit_tests/tasks/test_delete_conversation_task.py index d3158ffe4ee..52fc858fa57 100644 --- a/api/tests/unit_tests/tasks/test_delete_conversation_task.py +++ b/api/tests/unit_tests/tasks/test_delete_conversation_task.py @@ -208,7 +208,7 @@ def test_cleanup_removes_owned_resources_and_preserves_drive_files(sqlite_sessio key="drive.txt", file_kind=AgentDriveFileKind.TOOL_FILE, file_id=drive_file.id, - value_owned_by_drive=True, + value_owned_by_drive=False, is_skill=False, ), HumanInputFormRecipient( @@ -248,6 +248,11 @@ def test_cleanup_removes_owned_resources_and_preserves_drive_files(sqlite_sessio preserved_drive_file = sqlite_session.get(ToolFile, drive_file_id) assert preserved_drive_file is not None assert preserved_drive_file.conversation_id is None + preserved_drive_entry = sqlite_session.scalar( + select(AgentDriveFile).where(AgentDriveFile.file_id == drive_file_id) + ) + assert preserved_drive_entry is not None + assert preserved_drive_entry.value_owned_by_drive is True assert sqlite_session.get(ToolFile, other_file_id) is not None assert sqlite_session.get(Conversation, OTHER_CONVERSATION_ID) is not None