From dfac3e524e02ea2874114d333ce5dd193e488c4e Mon Sep 17 00:00:00 2001 From: zyssyz123 <916125788@qq.com> Date: Fri, 14 Aug 2026 14:30:04 +0000 Subject: [PATCH] fix: reclaim conversation resources on deletion (#40792) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/.env.example | 6 + api/configs/feature/__init__.py | 12 + api/extensions/ext_celery.py | 6 + ...4e050600_add_conversation_cleanup_index.py | 30 ++ api/models/model.py | 6 + api/services/agent_drive_service.py | 14 +- api/services/conversation_service.py | 9 +- api/tasks/delete_conversation_task.py | 190 ++++++++--- .../services/test_conversation_service.py | 11 +- .../tasks/test_delete_conversation_task.py | 178 ++++++++++ .../unit_tests/extensions/test_celery_ssl.py | 10 +- .../services/test_agent_drive_service.py | 28 +- .../services/test_conversation_service.py | 29 +- .../tasks/test_delete_conversation_task.py | 303 ++++++++++++++++++ docker/envs/core-services/shared.env.example | 5 + 15 files changed, 785 insertions(+), 52 deletions(-) create mode 100644 api/migrations/versions/2026_08_14_1200-56124e050600_add_conversation_cleanup_index.py create mode 100644 api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py create mode 100644 api/tests/unit_tests/tasks/test_delete_conversation_task.py diff --git a/api/.env.example b/api/.env.example index a11f6eb0d25..5bc4cffeb17 100644 --- a/api/.env.example +++ b/api/.env.example @@ -829,6 +829,12 @@ ENABLE_HUMAN_INPUT_TIMEOUT_TASK=true # Human input timeout check interval in minutes HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1 +# Whether to recover soft-deleted conversation cleanup jobs periodically +ENABLE_CONVERSATION_CLEANUP_TASK=true +# Recovery interval in minutes and maximum conversations dispatched per sweep +CONVERSATION_CLEANUP_TASK_INTERVAL=5 +CONVERSATION_CLEANUP_BATCH_SIZE=100 + # Nacos remote settings source HTTP timeouts (seconds). # Bound how long requests to the Nacos endpoint wait before failing, so a slow or # unresponsive Nacos server cannot stall API startup or token refresh. diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index ea2b4cacf80..0772b1d08d4 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1344,6 +1344,18 @@ class CeleryBeatConfig(BaseSettings): class CeleryScheduleTasksConfig(BaseSettings): + ENABLE_CONVERSATION_CLEANUP_TASK: bool = Field( + description="Enable periodic recovery of soft-deleted conversation cleanup", + default=True, + ) + CONVERSATION_CLEANUP_TASK_INTERVAL: PositiveInt = Field( + description="Soft-deleted conversation cleanup recovery interval in minutes", + default=5, + ) + CONVERSATION_CLEANUP_BATCH_SIZE: PositiveInt = Field( + description="Maximum soft-deleted conversations dispatched per cleanup sweep", + default=100, + ) ENABLE_CLEAN_EMBEDDING_CACHE_TASK: bool = Field( description="Enable clean embedding cache task", default=False, diff --git a/api/extensions/ext_celery.py b/api/extensions/ext_celery.py index a69971d7848..184dd5bddfa 100644 --- a/api/extensions/ext_celery.py +++ b/api/extensions/ext_celery.py @@ -181,6 +181,12 @@ def init_app(app: DifyApp) -> Celery: # if you add a new task, please add the switch to CeleryScheduleTasksConfig beat_schedule: dict[str, CeleryBeatScheduleEntry] = {} + if dify_config.ENABLE_CONVERSATION_CLEANUP_TASK: + imports.append("tasks.delete_conversation_task") + beat_schedule["conversation_cleanup_sweeper"] = { + "task": "tasks.delete_conversation_task.sweep_deleted_conversations", + "schedule": timedelta(minutes=dify_config.CONVERSATION_CLEANUP_TASK_INTERVAL), + } if dify_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK: imports.append("schedule.clean_embedding_cache_task") beat_schedule["clean_embedding_cache_task"] = { diff --git a/api/migrations/versions/2026_08_14_1200-56124e050600_add_conversation_cleanup_index.py b/api/migrations/versions/2026_08_14_1200-56124e050600_add_conversation_cleanup_index.py new file mode 100644 index 00000000000..13175fd6d63 --- /dev/null +++ b/api/migrations/versions/2026_08_14_1200-56124e050600_add_conversation_cleanup_index.py @@ -0,0 +1,30 @@ +"""add conversation cleanup index + +Revision ID: 56124e050600 +Revises: f3a9c2d17b4e +Create Date: 2026-08-14 12:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "56124e050600" +down_revision = "f3a9c2d17b4e" +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_index( + "conversation_is_deleted_updated_at_idx", + "conversations", + ["is_deleted", "updated_at"], + unique=False, + postgresql_where=sa.text("is_deleted IS true"), + ) + + +def downgrade(): + op.drop_index("conversation_is_deleted_updated_at_idx", table_name="conversations") diff --git a/api/models/model.py b/api/models/model.py index 8c02322d763..0cad0bb3feb 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1187,6 +1187,12 @@ class Conversation(Base): sa.text("updated_at DESC"), postgresql_where=sa.text("is_deleted IS false"), ), + sa.Index( + "conversation_is_deleted_updated_at_idx", + "is_deleted", + "updated_at", + postgresql_where=sa.text("is_deleted IS true"), + ), ) id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4())) diff --git a/api/services/agent_drive_service.py b/api/services/agent_drive_service.py index eb375f997d1..0e4722d1018 100644 --- a/api/services/agent_drive_service.py +++ b/api/services/agent_drive_service.py @@ -415,7 +415,12 @@ class AgentDriveService: file_kind = AgentDriveFileKind(item.file_ref.kind) file_id = item.file_ref.id size, mime_type, file_hash = self._validate_source( - session, tenant_id=tenant_id, user_id=user_id, file_kind=file_kind, file_id=file_id + session, + tenant_id=tenant_id, + user_id=user_id, + file_kind=file_kind, + file_id=file_id, + take_ownership=item.value_owned_by_drive, ) existing = session.scalar( @@ -725,6 +730,7 @@ class AgentDriveService: user_id: str, file_kind: AgentDriveFileKind, file_id: str, + take_ownership: bool = False, ) -> tuple[int | None, str | None, str | None]: """Verify the source file exists for the tenant (and user, for ToolFile). @@ -734,16 +740,20 @@ class AgentDriveService: try: if file_kind == AgentDriveFileKind.TOOL_FILE: tool_file = session.scalar( - select(ToolFile).where( + select(ToolFile) + .where( ToolFile.id == file_id, ToolFile.tenant_id == tenant_id, ToolFile.user_id == user_id, ) + .with_for_update() ) if tool_file is None: raise AgentDriveError( "source_not_found", "source ToolFile not found for this tenant/user", status_code=404 ) + if take_ownership: + tool_file.conversation_id = None return tool_file.size, tool_file.mimetype, None upload_file = session.scalar( select(UploadFile).where(UploadFile.id == file_id, UploadFile.tenant_id == tenant_id) diff --git a/api/services/conversation_service.py b/api/services/conversation_service.py index 406af068194..e4622056d28 100644 --- a/api/services/conversation_service.py +++ b/api/services/conversation_service.py @@ -228,7 +228,7 @@ class ConversationService: ) if retired_binding_id is None: raise AgentWorkspaceNotFoundError("Conversation participant Binding is unavailable") - session.delete(conversation) + conversation.is_deleted = True session.commit() except Exception: session.rollback() @@ -238,7 +238,12 @@ class ConversationService: tenant_id=app_model.tenant_id, binding_ids=(retired_binding_id,), ) - delete_conversation_related_data.delay(conversation.id) + try: + delete_conversation_related_data.delay(conversation.id) + except Exception: + # The soft-deleted row is a durable cleanup marker picked up by the + # periodic sweeper, so a broker outage must not resurrect or expose it. + logger.exception("Failed to enqueue cleanup for conversation %s", conversation.id) @classmethod def get_conversational_variable( diff --git a/api/tasks/delete_conversation_task.py b/api/tasks/delete_conversation_task.py index 6f426fe7fb0..c9bf4ce9f17 100644 --- a/api/tasks/delete_conversation_task.py +++ b/api/tasks/delete_conversation_task.py @@ -3,19 +3,127 @@ import time import click from celery import shared_task -from sqlalchemy import delete +from sqlalchemy import delete, select +from configs import dify_config from core.db.session_factory import session_factory -from models import ConversationVariable -from models.model import Message, MessageAnnotation, MessageFeedback +from extensions.ext_storage import storage +from models import ( + AgentDebugConversation, + Conversation, + ConversationVariable, + HumanInputForm, + HumanInputFormUploadFile, + HumanInputFormUploadToken, + Message, + MessageAgentThought, + MessageAnnotation, + MessageChain, + MessageFeedback, + MessageFile, + PinnedConversation, + SavedMessage, +) +from models.agent import AgentDriveFile, AgentDriveFileKind +from models.human_input import HumanInputDelivery, HumanInputFormRecipient from models.tools import ToolConversationVariables, ToolFile -from models.web import PinnedConversation logger = logging.getLogger(__name__) +_MAX_RETRIES = 5 +_RETRY_DELAY_SECONDS = 30 -@shared_task(queue="conversation") -def delete_conversation_related_data(conversation_id: str): + +def _delete_storage_object(file_key: str) -> None: + try: + storage.delete(file_key) + except Exception: + # A prior attempt may have deleted the object before its DB transaction + # rolled back. Only suppress the retry when the backend confirms absence. + if storage.exists(file_key): + raise + logger.info("Storage object %s was already absent", file_key) + + +def _cleanup_conversation_related_data(conversation_id: str) -> bool: + """Physically remove a soft-deleted conversation and its owned resources. + + 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 their Drive references + take over lifecycle ownership. + """ + + with session_factory.create_session() as session: + conversation = session.scalar(select(Conversation).where(Conversation.id == conversation_id).with_for_update()) + if conversation is not None and not conversation.is_deleted: + logger.warning("Skipped cleanup for active conversation %s", conversation_id) + return False + + tool_files = list( + session.scalars( + select(ToolFile) + .where(ToolFile.conversation_id == conversation_id) + .order_by(ToolFile.id) + .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) + + message_ids = select(Message.id).where(Message.conversation_id == conversation_id) + session.execute(delete(MessageAgentThought).where(MessageAgentThought.message_id.in_(message_ids))) + session.execute(delete(MessageChain).where(MessageChain.message_id.in_(message_ids))) + session.execute(delete(MessageFile).where(MessageFile.message_id.in_(message_ids))) + session.execute(delete(SavedMessage).where(SavedMessage.message_id.in_(message_ids))) + session.execute(delete(MessageAnnotation).where(MessageAnnotation.conversation_id == conversation_id)) + session.execute(delete(MessageFeedback).where(MessageFeedback.conversation_id == conversation_id)) + session.execute( + delete(ToolConversationVariables).where(ToolConversationVariables.conversation_id == conversation_id) + ) + session.execute(delete(ConversationVariable).where(ConversationVariable.conversation_id == conversation_id)) + + form_ids = select(HumanInputForm.id).where(HumanInputForm.conversation_id == conversation_id) + session.execute(delete(HumanInputFormUploadFile).where(HumanInputFormUploadFile.form_id.in_(form_ids))) + session.execute(delete(HumanInputFormUploadToken).where(HumanInputFormUploadToken.form_id.in_(form_ids))) + session.execute(delete(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id.in_(form_ids))) + session.execute(delete(HumanInputDelivery).where(HumanInputDelivery.form_id.in_(form_ids))) + session.execute(delete(HumanInputForm).where(HumanInputForm.conversation_id == conversation_id)) + + session.execute(delete(Message).where(Message.conversation_id == conversation_id)) + session.execute(delete(PinnedConversation).where(PinnedConversation.conversation_id == conversation_id)) + session.execute(delete(AgentDebugConversation).where(AgentDebugConversation.conversation_id == conversation_id)) + session.execute( + delete(Conversation).where( + Conversation.id == conversation_id, + Conversation.is_deleted.is_(True), + ) + ) + session.commit() + return True + + +@shared_task(queue="conversation", bind=True, max_retries=_MAX_RETRIES, default_retry_delay=_RETRY_DELAY_SECONDS) +def delete_conversation_related_data(self, conversation_id: str) -> None: """ Delete related data conversation in correct order from database to respect foreign key constraints @@ -28,38 +136,44 @@ def delete_conversation_related_data(conversation_id: str): ) start_at = time.perf_counter() + try: + cleaned = _cleanup_conversation_related_data(conversation_id) + except Exception as exc: + logger.exception("Failed to delete data for conversation_id: %s", conversation_id) + countdown = min(_RETRY_DELAY_SECONDS * (2**self.request.retries), 10 * 60) + raise self.retry(exc=exc, countdown=countdown) + + end_at = time.perf_counter() + logger.info( + click.style( + ( + f"Finished cleaning data for conversation_id {conversation_id}, " + f"cleaned={cleaned}, latency: {end_at - start_at}" + ), + fg="green", + ) + ) + + +@shared_task(queue="conversation") +def sweep_deleted_conversations() -> int: + """Re-enqueue soft-deleted conversations whose immediate dispatch was lost.""" + with session_factory.create_session() as session: + conversation_ids = list( + session.scalars( + select(Conversation.id) + .where(Conversation.is_deleted.is_(True)) + .order_by(Conversation.updated_at, Conversation.id) + .limit(dify_config.CONVERSATION_CLEANUP_BATCH_SIZE) + ) + ) + + dispatched = 0 + for conversation_id in conversation_ids: try: - session.execute(delete(MessageAnnotation).where(MessageAnnotation.conversation_id == conversation_id)) - - session.execute(delete(MessageFeedback).where(MessageFeedback.conversation_id == conversation_id)) - - session.execute( - delete(ToolConversationVariables).where(ToolConversationVariables.conversation_id == conversation_id) - ) - - session.execute(delete(ToolFile).where(ToolFile.conversation_id == conversation_id)) - - session.execute(delete(ConversationVariable).where(ConversationVariable.conversation_id == conversation_id)) - - session.execute(delete(Message).where(Message.conversation_id == conversation_id)) - - session.execute(delete(PinnedConversation).where(PinnedConversation.conversation_id == conversation_id)) - - session.commit() - - end_at = time.perf_counter() - logger.info( - click.style( - ( - f"Succeeded cleaning data from db for conversation_id {conversation_id} " - f"latency: {end_at - start_at}" - ), - fg="green", - ) - ) - + delete_conversation_related_data.delay(conversation_id) + dispatched += 1 except Exception: - logger.exception("Failed to delete data from db for conversation_id: %s failed", conversation_id) - session.rollback() - raise + logger.exception("Failed to dispatch cleanup for conversation %s", conversation_id) + return dispatched diff --git a/api/tests/test_containers_integration_tests/services/test_conversation_service.py b/api/tests/test_containers_integration_tests/services/test_conversation_service.py index 75f017fbc50..69407cb039b 100644 --- a/api/tests/test_containers_integration_tests/services/test_conversation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_conversation_service.py @@ -1081,7 +1081,7 @@ class TestConversationServiceExport: Test conversation deletion with async cleanup. Deletion is a two-step process: - 1. Immediately delete the conversation record from database + 1. Immediately hide the conversation with a durable soft-delete marker 2. Trigger async background task to clean up related data (messages, annotations, vector embeddings, file uploads) """ @@ -1102,9 +1102,10 @@ class TestConversationServiceExport: ) # Assert - Verify two-step deletion process - # Step 1: Immediate database deletion + # Step 1: Immediate logical deletion deleted = db_session_with_containers.scalar(select(Conversation).where(Conversation.id == conversation_id)) - assert deleted is None + assert deleted is not None + assert deleted.is_deleted is True # Step 2: Async cleanup task triggered # The Celery task will handle cleanup of messages, annotations, etc. @@ -1166,8 +1167,8 @@ class TestConversationServiceExport: ) conversation_id = conversation.id - # Act — force an error during the delete to exercise the rollback path - with patch.object(db_session_with_containers, "delete", side_effect=Exception("DB error")): + # Act — force an error during the soft-delete commit to exercise rollback + with patch.object(db_session_with_containers, "commit", side_effect=Exception("DB error")): with pytest.raises(Exception, match="DB error"): ConversationService.delete( app_model=app_model, 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 new file mode 100644 index 00000000000..b7b9e934cc4 --- /dev/null +++ b/api/tests/test_containers_integration_tests/tasks/test_delete_conversation_task.py @@ -0,0 +1,178 @@ +from threading import Event, Thread +from unittest.mock import patch + +from sqlalchemy import event, select +from sqlalchemy.orm import Session + +from models import AppMode, Conversation, ToolFile +from models.agent import AgentDriveFile, AgentDriveFileKind +from models.enums import ConversationFromSource, ConversationStatus +from tasks.delete_conversation_task import _cleanup_conversation_related_data + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +APP_ID = "22222222-2222-2222-2222-222222222222" +ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" +CONVERSATION_ID = "44444444-4444-4444-4444-444444444444" +AGENT_ID = "55555555-5555-5555-5555-555555555555" + + +def test_cleanup_deletes_owned_storage_and_preserves_drive_file( + 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, + ) + owned_file = ToolFile( + user_id=ACCOUNT_ID, + tenant_id=TENANT_ID, + conversation_id=CONVERSATION_ID, + file_key=f"tools/{TENANT_ID}/owned.txt", + mimetype="text/plain", + name="owned.txt", + size=5, + ) + drive_file = ToolFile( + user_id=ACCOUNT_ID, + tenant_id=TENANT_ID, + conversation_id=CONVERSATION_ID, + file_key=f"tools/{TENANT_ID}/drive.txt", + mimetype="text/plain", + name="drive.txt", + size=5, + ) + db_session_with_containers.add_all([conversation, owned_file, drive_file]) + db_session_with_containers.flush() + drive_entry = AgentDriveFile( + tenant_id=TENANT_ID, + agent_id=AGENT_ID, + key="drive.txt", + file_kind=AgentDriveFileKind.TOOL_FILE, + file_id=drive_file.id, + value_owned_by_drive=False, + is_skill=False, + ) + db_session_with_containers.add(drive_entry) + db_session_with_containers.commit() + owned_file_id = owned_file.id + drive_file_id = drive_file.id + + with patch("tasks.delete_conversation_task.storage") as storage_mock: + assert _cleanup_conversation_related_data(CONVERSATION_ID) is True + + storage_mock.delete.assert_called_once_with(f"tools/{TENANT_ID}/owned.txt") + db_session_with_containers.expire_all() + assert db_session_with_containers.get(Conversation, CONVERSATION_ID) is None + assert db_session_with_containers.get(ToolFile, owned_file_id) is None + 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 + + +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() + + 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: + 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/extensions/test_celery_ssl.py b/api/tests/unit_tests/extensions/test_celery_ssl.py index ad8192ecea4..c5df0a3b7ff 100644 --- a/api/tests/unit_tests/extensions/test_celery_ssl.py +++ b/api/tests/unit_tests/extensions/test_celery_ssl.py @@ -162,6 +162,8 @@ class TestCelerySSLConfiguration: # Mock all the scheduler configs mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False + mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False @@ -195,7 +197,7 @@ class TestCelerySSLConfiguration: assert "redis_backend_use_ssl" in celery_app.conf assert celery_app.conf["redis_backend_use_ssl"] is not None - def test_celery_init_applies_global_keyprefix_and_registers_agent_resource_collector(self): + def test_celery_init_registers_required_agent_and_conversation_tasks(self): mock_config = MagicMock() mock_config.BROKER_USE_SSL = False mock_config.REDIS_KEY_PREFIX = "enterprise-a" @@ -210,6 +212,8 @@ class TestCelerySSLConfiguration: mock_config.CELERY_TASK_ANNOTATIONS = {} mock_config.CELERY_BEAT_SCHEDULER_TIME = 1 + mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True + mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5 mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False mock_config.ENABLE_CLEAN_UNUSED_DATASETS_TASK = False mock_config.ENABLE_CREATE_TIDB_SERVERLESS_TASK = False @@ -241,3 +245,7 @@ class TestCelerySSLConfiguration: assert celery_app.conf["broker_transport_options"]["global_keyprefix"] == "enterprise-a:" assert celery_app.conf["result_backend_transport_options"]["global_keyprefix"] == "enterprise-a:" assert "tasks.collect_agent_resources_task" in celery_app.conf["imports"] + assert "tasks.delete_conversation_task" in celery_app.conf["imports"] + assert celery_app.conf["beat_schedule"]["conversation_cleanup_sweeper"]["task"] == ( + "tasks.delete_conversation_task.sweep_deleted_conversations" + ) diff --git a/api/tests/unit_tests/services/test_agent_drive_service.py b/api/tests/unit_tests/services/test_agent_drive_service.py index 74ee1a5d2f1..e917745b4e7 100644 --- a/api/tests/unit_tests/services/test_agent_drive_service.py +++ b/api/tests/unit_tests/services/test_agent_drive_service.py @@ -91,11 +91,11 @@ def _seed_agent(*, tenant_id: str = TENANT, agent_id: str = AGENT) -> None: session.commit() -def _seed_tool_file(*, user_id: str = USER, name: str = "f.txt") -> str: +def _seed_tool_file(*, user_id: str = USER, name: str = "f.txt", conversation_id: str | None = None) -> str: tool_file = ToolFile( user_id=user_id, tenant_id=TENANT, - conversation_id=None, + conversation_id=conversation_id, file_key=f"tools/{TENANT}/{name}", mimetype="text/plain", name=name, @@ -156,6 +156,30 @@ def test_commit_then_manifest_lists_the_entry(): ) +def test_commit_owned_tool_file_detaches_conversation_ownership(): + conversation_id = "44444444-4444-4444-4444-444444444444" + tool_file_id = _seed_tool_file(conversation_id=conversation_id) + + _commit("data/report.txt", tool_file_id, owned=True) + + with session_factory.create_session() as session: + tool_file = session.get(ToolFile, tool_file_id) + assert tool_file is not None + assert tool_file.conversation_id is None + + +def test_commit_shared_tool_file_keeps_conversation_ownership(): + conversation_id = "44444444-4444-4444-4444-444444444444" + tool_file_id = _seed_tool_file(conversation_id=conversation_id) + + _commit("data/report.txt", tool_file_id, owned=False) + + with session_factory.create_session() as session: + tool_file = session.get(ToolFile, tool_file_id) + assert tool_file is not None + assert tool_file.conversation_id == conversation_id + + def test_commit_skill_row_persists_metadata_and_lists_catalog() -> None: tf = _seed_tool_file(name="SKILL.md") AgentDriveService().commit( diff --git a/api/tests/unit_tests/services/test_conversation_service.py b/api/tests/unit_tests/services/test_conversation_service.py index 3d73922fd4e..9476de34394 100644 --- a/api/tests/unit_tests/services/test_conversation_service.py +++ b/api/tests/unit_tests/services/test_conversation_service.py @@ -188,13 +188,16 @@ def test_delete_retires_then_commits_before_enqueue(monkeypatch: pytest.MonkeyPa "enqueue_agent_resource_collection", MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), ) - monkeypatch.setattr(conversation_service.delete_conversation_related_data, "delay", MagicMock()) + delete_related = MagicMock() + monkeypatch.setattr(conversation_service.delete_conversation_related_data, "delay", delete_related) ConversationService.delete(app, conversation.id, None, session=sqlite_session) assert events == ["retire", "commit", "enqueue"] + assert conversation.is_deleted is True assert get_binding.call_args.kwargs["binding_id"] == "conversation-binding-1" assert retire_binding.call_args.kwargs["binding_id"] == "conversation-binding-1" + delete_related.assert_called_once_with(conversation.id) def test_delete_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: @@ -202,7 +205,7 @@ def test_delete_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch, conversation = ConversationServiceTestDataFactory.create_conversation() conversation.agent_workspace_binding_id = "binding-1" sqlite_session.add(conversation) - sqlite_session.flush() + sqlite_session.commit() rollback_events: list[str] = [] event.listen(sqlite_session, "after_rollback", lambda _session: rollback_events.append("rollback")) @@ -226,10 +229,32 @@ def test_delete_commit_failure_does_not_enqueue(monkeypatch: pytest.MonkeyPatch, ConversationService.delete(app, conversation.id, None, session=sqlite_session) assert rollback_events == ["rollback"] + assert conversation.is_deleted is False enqueue_collection.assert_not_called() delete_related.assert_not_called() +def test_delete_keeps_soft_deleted_marker_when_dispatch_fails( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + app = ConversationServiceTestDataFactory.create_app() + conversation = ConversationServiceTestDataFactory.create_conversation() + sqlite_session.add(conversation) + sqlite_session.flush() + monkeypatch.setattr(ConversationService, "get_conversation", MagicMock(return_value=conversation)) + monkeypatch.setattr( + conversation_service.delete_conversation_related_data, + "delay", + MagicMock(side_effect=RuntimeError("broker unavailable")), + ) + + ConversationService.delete(app, conversation.id, None, session=sqlite_session) + + persisted = sqlite_session.get(Conversation, conversation.id) + assert persisted is not None + assert persisted.is_deleted is True + + class TestConversationServicePagination: """Test conversation pagination operations.""" diff --git a/api/tests/unit_tests/tasks/test_delete_conversation_task.py b/api/tests/unit_tests/tasks/test_delete_conversation_task.py new file mode 100644 index 00000000000..d928549975b --- /dev/null +++ b/api/tests/unit_tests/tasks/test_delete_conversation_task.py @@ -0,0 +1,303 @@ +from datetime import UTC, datetime, timedelta +from decimal import Decimal +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy import select +from sqlalchemy.orm import Session + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.workflow.human_input_adapter import DeliveryMethodType +from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus +from graphon.file import FileTransferMethod, FileType +from models import ( + AgentDebugConversation, + AppMode, + Conversation, + ConversationVariable, + HumanInputForm, + HumanInputFormUploadFile, + HumanInputFormUploadToken, + Message, + MessageAgentThought, + MessageAnnotation, + MessageChain, + MessageFeedback, + MessageFile, + PinnedConversation, + SavedMessage, +) +from models.agent import AgentConfigDraftType, AgentDriveFile, AgentDriveFileKind +from models.enums import ( + ConversationFromSource, + ConversationStatus, + CreatorUserRole, + FeedbackFromSource, + FeedbackRating, + MessageChainType, +) +from models.human_input import HumanInputDelivery, HumanInputFormRecipient, RecipientType +from models.tools import ToolConversationVariables, ToolFile +from tasks.delete_conversation_task import _cleanup_conversation_related_data, sweep_deleted_conversations + +TENANT_ID = "11111111-1111-1111-1111-111111111111" +APP_ID = "22222222-2222-2222-2222-222222222222" +ACCOUNT_ID = "33333333-3333-3333-3333-333333333333" +CONVERSATION_ID = "44444444-4444-4444-4444-444444444444" +OTHER_CONVERSATION_ID = "55555555-5555-5555-5555-555555555555" +MESSAGE_ID = "66666666-6666-6666-6666-666666666666" +AGENT_ID = "77777777-7777-7777-7777-777777777777" + + +def _conversation(conversation_id: str, *, deleted: bool) -> Conversation: + return Conversation( + id=conversation_id, + app_id=APP_ID, + mode=AppMode.CHAT, + name="Test conversation", + inputs={}, + status=ConversationStatus.NORMAL, + from_source=ConversationFromSource.CONSOLE, + from_account_id=ACCOUNT_ID, + is_deleted=deleted, + ) + + +def _message() -> Message: + return Message( + id=MESSAGE_ID, + app_id=APP_ID, + conversation_id=CONVERSATION_ID, + inputs={}, + query="hello", + message={"role": "user", "content": "hello"}, + answer="world", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + invoke_from=InvokeFrom.WEB_APP, + from_source=ConversationFromSource.CONSOLE, + from_account_id=ACCOUNT_ID, + ) + + +def _tool_file(*, name: str, conversation_id: str | None = CONVERSATION_ID) -> ToolFile: + return ToolFile( + user_id=ACCOUNT_ID, + tenant_id=TENANT_ID, + conversation_id=conversation_id, + file_key=f"tools/{TENANT_ID}/{name}", + mimetype="text/plain", + name=name, + size=5, + ) + + +def test_cleanup_removes_owned_resources_and_preserves_drive_files(sqlite_session: Session) -> None: + conversation = _conversation(CONVERSATION_ID, deleted=True) + other_conversation = _conversation(OTHER_CONVERSATION_ID, deleted=False) + message = _message() + owned_file = _tool_file(name="owned.txt") + drive_file = _tool_file(name="drive.txt") + other_file = _tool_file(name="other.txt", conversation_id=OTHER_CONVERSATION_ID) + sqlite_session.add_all([conversation, other_conversation, message, owned_file, drive_file, other_file]) + sqlite_session.flush() + + message_chain = MessageChain(message_id=MESSAGE_ID, type=MessageChainType.SYSTEM, input=None, output=None) + form = HumanInputForm( + tenant_id=TENANT_ID, + app_id=APP_ID, + workflow_run_id=None, + conversation_id=CONVERSATION_ID, + form_kind=HumanInputFormKind.RUNTIME, + node_id="ask-human", + form_definition="{}", + rendered_content="form", + status=HumanInputFormStatus.WAITING, + expiration_time=datetime.now(UTC) + timedelta(hours=1), + ) + sqlite_session.add_all([message_chain, form]) + sqlite_session.flush() + + delivery = HumanInputDelivery( + form_id=form.id, + delivery_method_type=DeliveryMethodType.WEBAPP, + delivery_config_id=None, + channel_payload="{}", + ) + sqlite_session.add(delivery) + sqlite_session.flush() + + upload_token = HumanInputFormUploadToken( + tenant_id=TENANT_ID, + app_id=APP_ID, + form_id=form.id, + recipient_id="88888888-8888-8888-8888-888888888888", + token="upload-token", + ) + sqlite_session.add(upload_token) + sqlite_session.flush() + + related_rows = [ + MessageAgentThought( + message_id=MESSAGE_ID, + position=1, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=ACCOUNT_ID, + message_chain_id=message_chain.id, + ), + MessageFile( + message_id=MESSAGE_ID, + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.REMOTE_URL, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=ACCOUNT_ID, + url="https://example.com/file.txt", + ), + SavedMessage( + app_id=APP_ID, + message_id=MESSAGE_ID, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=ACCOUNT_ID, + ), + MessageAnnotation( + app_id=APP_ID, + question="hello", + content="world", + account_id=ACCOUNT_ID, + conversation_id=CONVERSATION_ID, + message_id=MESSAGE_ID, + ), + MessageFeedback( + app_id=APP_ID, + conversation_id=CONVERSATION_ID, + message_id=MESSAGE_ID, + rating=FeedbackRating.LIKE, + from_source=FeedbackFromSource.ADMIN, + from_account_id=ACCOUNT_ID, + ), + ToolConversationVariables( + user_id=ACCOUNT_ID, + tenant_id=TENANT_ID, + conversation_id=CONVERSATION_ID, + variables_str="{}", + ), + ConversationVariable( + id="99999999-9999-9999-9999-999999999999", + conversation_id=CONVERSATION_ID, + app_id=APP_ID, + data="{}", + ), + PinnedConversation( + app_id=APP_ID, + conversation_id=CONVERSATION_ID, + created_by_role=CreatorUserRole.ACCOUNT, + created_by=ACCOUNT_ID, + ), + AgentDebugConversation( + tenant_id=TENANT_ID, + agent_id=AGENT_ID, + app_id=APP_ID, + account_id=ACCOUNT_ID, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + conversation_id=CONVERSATION_ID, + ), + AgentDriveFile( + tenant_id=TENANT_ID, + agent_id=AGENT_ID, + key="drive.txt", + file_kind=AgentDriveFileKind.TOOL_FILE, + file_id=drive_file.id, + value_owned_by_drive=False, + is_skill=False, + ), + HumanInputFormRecipient( + form_id=form.id, + delivery_id=delivery.id, + recipient_type=RecipientType.CONSOLE, + recipient_payload="{}", + ), + HumanInputFormUploadFile( + tenant_id=TENANT_ID, + app_id=APP_ID, + form_id=form.id, + upload_file_id="aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa", + upload_token_id=upload_token.id, + ), + ] + sqlite_session.add_all(related_rows) + sqlite_session.commit() + form_id = form.id + owned_file_id = owned_file.id + owned_file_key = owned_file.file_key + drive_file_id = drive_file.id + other_file_id = other_file.id + + with patch("tasks.delete_conversation_task.storage") as storage_mock: + assert _cleanup_conversation_related_data(CONVERSATION_ID) is True + + storage_mock.delete.assert_called_once_with(owned_file_key) + sqlite_session.expire_all() + assert sqlite_session.get(Conversation, CONVERSATION_ID) is None + assert sqlite_session.get(Message, MESSAGE_ID) is None + assert ( + sqlite_session.scalar(select(MessageAgentThought).where(MessageAgentThought.message_id == MESSAGE_ID)) is None + ) + assert sqlite_session.scalar(select(HumanInputForm).where(HumanInputForm.id == form_id)) is None + assert sqlite_session.get(ToolFile, owned_file_id) is None + 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 + + +def test_cleanup_storage_failure_retains_marker_and_file_key(sqlite_session: Session) -> None: + conversation = _conversation(CONVERSATION_ID, deleted=True) + tool_file = _tool_file(name="retry.txt") + sqlite_session.add_all([conversation, tool_file]) + sqlite_session.commit() + + with patch("tasks.delete_conversation_task.storage") as storage_mock: + storage_mock.delete.side_effect = RuntimeError("storage unavailable") + storage_mock.exists.return_value = True + with pytest.raises(RuntimeError, match="storage unavailable"): + _cleanup_conversation_related_data(CONVERSATION_ID) + + sqlite_session.expire_all() + persisted_conversation = sqlite_session.get(Conversation, CONVERSATION_ID) + assert persisted_conversation is not None + assert persisted_conversation.is_deleted is True + assert sqlite_session.get(ToolFile, tool_file.id) is not None + + +def test_cleanup_skips_active_conversation(sqlite_session: Session) -> None: + conversation = _conversation(CONVERSATION_ID, deleted=False) + tool_file = _tool_file(name="active.txt") + sqlite_session.add_all([conversation, tool_file]) + sqlite_session.commit() + + with patch("tasks.delete_conversation_task.storage") as storage_mock: + assert _cleanup_conversation_related_data(CONVERSATION_ID) is False + storage_mock.delete.assert_not_called() + + assert sqlite_session.get(Conversation, CONVERSATION_ID) is not None + assert sqlite_session.get(ToolFile, tool_file.id) is not None + + +def test_sweeper_dispatches_only_soft_deleted_conversations(sqlite_session: Session) -> None: + sqlite_session.add_all( + [ + _conversation(CONVERSATION_ID, deleted=True), + _conversation(OTHER_CONVERSATION_ID, deleted=False), + ] + ) + sqlite_session.commit() + + with patch("tasks.delete_conversation_task.delete_conversation_related_data.delay", MagicMock()) as delay: + assert sweep_deleted_conversations.run() == 1 + + delay.assert_called_once_with(CONVERSATION_ID) diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index ed34a894809..f4c1211e6b4 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -511,6 +511,11 @@ MILVUS_ENABLE_HYBRID_SEARCH=False ENABLE_HUMAN_INPUT_TIMEOUT_TASK=true HUMAN_INPUT_TIMEOUT_TASK_INTERVAL=1 +# Conversation cleanup recovery task +ENABLE_CONVERSATION_CLEANUP_TASK=true +CONVERSATION_CLEANUP_TASK_INTERVAL=5 +CONVERSATION_CLEANUP_BATCH_SIZE=100 + # Nacos remote settings source HTTP timeouts (seconds). # Bound how long requests to the Nacos endpoint wait before failing, so a slow or # unresponsive Nacos server cannot stall API startup or token refresh.