fix: serialize conversation cleanup with drive commits

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)
 

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

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