fix: reclaim resources after conversation deletion

This commit is contained in:
Yansong Zhang 2026-08-14 19:41:47 +08:00
parent ab41571b28
commit f600afe2ca
15 changed files with 675 additions and 52 deletions

View File

@ -830,6 +830,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.

View File

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

View File

@ -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"] = {

View File

@ -0,0 +1,30 @@
"""add conversation cleanup index
Revision ID: 56124e050600
Revises: a1c7f4e9b3d2
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 = "a1c7f4e9b3d2"
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")

View File

@ -1183,6 +1183,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()))

View File

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

View File

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

View File

@ -3,19 +3,120 @@ import time
import click
from celery import shared_task
from sqlalchemy import delete
from sqlalchemy import delete, select, update
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 survive this lifecycle.
"""
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
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)
.where(ToolFile.conversation_id == conversation_id)
.order_by(ToolFile.id)
.with_for_update()
)
)
for tool_file in tool_files:
_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 +129,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

View File

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

View File

@ -0,0 +1,78 @@
from unittest.mock import patch
from sqlalchemy import 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=True,
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
assert db_session_with_containers.scalar(
select(AgentDriveFile).where(AgentDriveFile.file_id == drive_file_id)
) is not None

View File

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

View File

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

View File

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

View File

@ -0,0 +1,300 @@
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=True,
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
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)

View File

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