Optimize expired message cleanup batching

api/configs/feature/__init__.py

@@ -1460,6 +1460,18 @@ class SandboxExpiredRecordsCleanConfig(BaseSettings):
         description="Maximum number of records to process in each batch",
         default=1000,
     )
+    SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_SIZE: PositiveInt = Field(
+        description="Initial number of message candidates to scan in each expired records cleanup batch",
+        default=10000,
+    )
+    SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_MAX_SIZE: PositiveInt = Field(
+        description="Maximum number of message candidates to scan in each expired records cleanup batch",
+        default=50000,
+    )
+    SANDBOX_EXPIRED_RECORDS_CLEAN_DELETE_BATCH_SIZE: PositiveInt = Field(
+        description="Maximum number of records to delete in each expired records cleanup transaction",
+        default=1000,
+    )
     SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL: PositiveInt = Field(
         description="Maximum interval in milliseconds between batches",
         default=200,

api/schedule/clean_messages.py

@@ -40,7 +40,9 @@ def clean_messages():
             service = MessagesCleanService.from_days(
                 policy=policy,
                 days=dify_config.SANDBOX_EXPIRED_RECORDS_RETENTION_DAYS,
-                batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_SIZE,
+                batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_SIZE,
+                max_candidate_batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_CANDIDATE_BATCH_MAX_SIZE,
+                delete_batch_size=dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_DELETE_BATCH_SIZE,
             )
             stats = service.run()
 

api/services/retention/conversation/messages_clean_policy.py

@@ -70,6 +70,13 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
     - Safe default: if tenant mapping or plan is missing, do NOT delete
     """
 
+    _graceful_period_days: int
+    _tenant_whitelist: Sequence[str]
+    _tenant_whitelist_set: set[str]
+    _plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]]
+    _current_timestamp: int | None
+    _plan_cache: dict[str, SubscriptionPlan | None]
+
     def __init__(
         self,
         plan_provider: Callable[[Sequence[str]], dict[str, SubscriptionPlan]],
@@ -79,8 +86,10 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
     ) -> None:
         self._graceful_period_days = graceful_period_days
         self._tenant_whitelist: Sequence[str] = tenant_whitelist or []
+        self._tenant_whitelist_set = set(self._tenant_whitelist)
         self._plan_provider = plan_provider
         self._current_timestamp = current_timestamp
+        self._plan_cache = {}
 
     @override
     def filter_message_ids(
@@ -101,9 +110,12 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
         if not messages or not app_to_tenant:
             return []
 
-        # Get unique tenant_ids and fetch subscription plans
-        tenant_ids = list(set(app_to_tenant.values()))
-        tenant_plans = self._plan_provider(tenant_ids)
+        # Get unique tenant_ids and fetch subscription plans. Plans are cached for the whole
+        # policy lifetime because message cleanup evaluates many adjacent batches from the same apps.
+        tenant_ids = sorted(
+            {tenant_id for tenant_id in app_to_tenant.values() if tenant_id not in self._tenant_whitelist_set}
+        )
+        tenant_plans = self._get_tenant_plans(tenant_ids)
 
         if not tenant_plans:
             return []
@@ -115,6 +127,28 @@ class BillingSandboxPolicy(MessagesCleanPolicy):
             tenant_plans=tenant_plans,
         )
 
+    def _get_tenant_plans(self, tenant_ids: Sequence[str]) -> dict[str, SubscriptionPlan]:
+        """
+        Return cached subscription plans for tenant ids.
+
+        Missing billing responses are cached as None and remain a safe non-delete decision.
+        Provider exceptions still propagate so transient billing failures do not silently change cleanup behavior.
+        """
+        unique_tenant_ids = sorted(set(tenant_ids))
+        missing_tenant_ids = [tenant_id for tenant_id in unique_tenant_ids if tenant_id not in self._plan_cache]
+
+        if missing_tenant_ids:
+            fetched_plans = self._plan_provider(missing_tenant_ids)
+            for tenant_id in missing_tenant_ids:
+                self._plan_cache[tenant_id] = fetched_plans.get(tenant_id)
+
+        plans: dict[str, SubscriptionPlan] = {}
+        for tenant_id in unique_tenant_ids:
+            plan = self._plan_cache.get(tenant_id)
+            if plan is not None:
+                plans[tenant_id] = plan
+        return plans
+
     def _filter_expired_sandbox_messages(
         self,
         messages: Sequence[SimpleMessage],

api/services/retention/conversation/messages_clean_service.py

@@ -1,8 +1,9 @@
 import datetime
 import logging
+import math
 import random
 import time
-from collections.abc import Sequence
+from collections.abc import Iterator, Sequence
 from typing import TYPE_CHECKING, TypedDict, cast
 
 import sqlalchemy as sa
@@ -37,6 +38,10 @@ if TYPE_CHECKING:
     from opentelemetry.metrics import Counter, Histogram
 
 
+_MIN_ELIGIBLE_HIT_RATE = 0.005
+_HIT_RATE_EMA_ALPHA = 0.3
+
+
 class MessagesCleanupMetrics:
     """
     Records low-cardinality OpenTelemetry metrics for expired message cleanup jobs.
@@ -165,12 +170,24 @@ class MessagesCleanStatsDict(TypedDict):
     total_deleted: int
 
 
+class MessageDeleteResultDict(TypedDict):
+    messages_deleted: int
+    chunks: int
+    relations_ms: int
+    messages_ms: int
+
+
 class MessagesCleanService:
     """
     Service for cleaning expired messages based on retention policies.
 
     Compatible with non cloud edition (billing disabled): all messages in the time range will be deleted.
     If billing is enabled: only sandbox plan tenant messages are deleted (with whitelist and grace period support).
+
+    The scan cursor advances by candidate messages, not messages selected by the policy. This keeps cleanup moving
+    through windows dominated by paid, unknown, or grace-period tenants without asking SQL to find enough eligible rows.
+    Candidate scan batches can grow independently from delete batches, while deletes stay chunked into
+    small transactions.
     """
 
     def __init__(
@@ -179,6 +196,8 @@ class MessagesCleanService:
         end_before: datetime.datetime,
         start_from: datetime.datetime | None = None,
         batch_size: int = 1000,
+        max_candidate_batch_size: int | None = None,
+        delete_batch_size: int | None = None,
         dry_run: bool = False,
         task_label: str = "custom",
     ) -> None:
@@ -189,14 +208,28 @@ class MessagesCleanService:
             policy: The policy that determines which messages to delete
             end_before: End time (exclusive) of the range
             start_from: Optional start time (inclusive) of the range
-            batch_size: Number of messages to process per batch
+            batch_size: Initial number of candidate messages to scan per batch
+            max_candidate_batch_size: Maximum number of candidate messages to scan per batch
+            delete_batch_size: Maximum number of messages to delete per transaction
             dry_run: Whether to perform a dry run (no actual deletion)
             task_label: Optional task label for retention metrics
         """
+        if batch_size <= 0:
+            raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
+
+        if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
+            raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
+
+        if delete_batch_size is not None and delete_batch_size <= 0:
+            raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
+
         self._policy = policy
         self._end_before = end_before
         self._start_from = start_from
         self._batch_size = batch_size
+        self._candidate_batch_size = batch_size
+        self._delete_batch_size = delete_batch_size or batch_size
+        self._max_candidate_batch_size = max(max_candidate_batch_size or batch_size, self._candidate_batch_size)
         self._dry_run = dry_run
         self._metrics = MessagesCleanupMetrics(
             dry_run=dry_run,
@@ -211,6 +244,8 @@ class MessagesCleanService:
         start_from: datetime.datetime,
         end_before: datetime.datetime,
         batch_size: int = 1000,
+        max_candidate_batch_size: int | None = None,
+        delete_batch_size: int | None = None,
         dry_run: bool = False,
         task_label: str = "custom",
     ) -> "MessagesCleanService":
@@ -223,7 +258,9 @@ class MessagesCleanService:
             policy: The policy that determines which messages to delete
             start_from: Start time (inclusive) of the range
             end_before: End time (exclusive) of the range
-            batch_size: Number of messages to process per batch
+            batch_size: Initial number of candidate messages to scan per batch
+            max_candidate_batch_size: Maximum number of candidate messages to scan per batch
+            delete_batch_size: Maximum number of messages to delete per transaction
             dry_run: Whether to perform a dry run (no actual deletion)
             task_label: Optional task label for retention metrics
 
@@ -239,11 +276,20 @@ class MessagesCleanService:
         if batch_size <= 0:
             raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
 
+        if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
+            raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
+
+        if delete_batch_size is not None and delete_batch_size <= 0:
+            raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
+
         logger.info(
-            "clean_messages: start_from=%s, end_before=%s, batch_size=%s, policy=%s",
+            "clean_messages: start_from=%s, end_before=%s, batch_size=%s, "
+            "max_candidate_batch_size=%s, delete_batch_size=%s, policy=%s",
             start_from,
             end_before,
             batch_size,
+            max_candidate_batch_size,
+            delete_batch_size,
             policy.__class__.__name__,
         )
 
@@ -252,6 +298,8 @@ class MessagesCleanService:
             end_before=end_before,
             start_from=start_from,
             batch_size=batch_size,
+            max_candidate_batch_size=max_candidate_batch_size,
+            delete_batch_size=delete_batch_size,
             dry_run=dry_run,
             task_label=task_label,
         )
@@ -262,6 +310,8 @@ class MessagesCleanService:
         policy: MessagesCleanPolicy,
         days: int = 30,
         batch_size: int = 1000,
+        max_candidate_batch_size: int | None = None,
+        delete_batch_size: int | None = None,
         dry_run: bool = False,
         task_label: str = "custom",
     ) -> "MessagesCleanService":
@@ -271,7 +321,9 @@ class MessagesCleanService:
         Args:
             policy: The policy that determines which messages to delete
             days: Number of days to look back from now
-            batch_size: Number of messages to process per batch
+            batch_size: Initial number of candidate messages to scan per batch
+            max_candidate_batch_size: Maximum number of candidate messages to scan per batch
+            delete_batch_size: Maximum number of messages to delete per transaction
             dry_run: Whether to perform a dry run (no actual deletion)
             task_label: Optional task label for retention metrics
 
@@ -287,13 +339,22 @@ class MessagesCleanService:
         if batch_size <= 0:
             raise ValueError(f"batch_size ({batch_size}) must be greater than 0")
 
+        if max_candidate_batch_size is not None and max_candidate_batch_size <= 0:
+            raise ValueError(f"max_candidate_batch_size ({max_candidate_batch_size}) must be greater than 0")
+
+        if delete_batch_size is not None and delete_batch_size <= 0:
+            raise ValueError(f"delete_batch_size ({delete_batch_size}) must be greater than 0")
+
         end_before = naive_utc_now() - datetime.timedelta(days=days)
 
         logger.info(
-            "clean_messages: days=%s, end_before=%s, batch_size=%s, policy=%s",
+            "clean_messages: days=%s, end_before=%s, batch_size=%s, "
+            "max_candidate_batch_size=%s, delete_batch_size=%s, policy=%s",
             days,
             end_before,
             batch_size,
+            max_candidate_batch_size,
+            delete_batch_size,
             policy.__class__.__name__,
         )
 
@@ -302,6 +363,8 @@ class MessagesCleanService:
             end_before=end_before,
             start_from=None,
             batch_size=batch_size,
+            max_candidate_batch_size=max_candidate_batch_size,
+            delete_batch_size=delete_batch_size,
             dry_run=dry_run,
             task_label=task_label,
         )
@@ -333,10 +396,10 @@ class MessagesCleanService:
         Time range is [start_from, end_before)
 
         Steps:
-        1. Iterate messages using cursor pagination (by created_at, id)
-        2. Query app_id -> tenant_id mapping
+        1. Iterate candidate messages using cursor pagination (by created_at, id)
+        2. Resolve app_id -> tenant_id mapping with a job-level cache
         3. Delegate to policy to determine which messages to delete
-        4. Batch delete messages and their relations
+        4. Delete messages and their relations in small chunks
 
         Returns:
             Dict with statistics: batches, filtered_messages, total_deleted
@@ -348,19 +411,23 @@ class MessagesCleanService:
             "total_deleted": 0,
         }
 
-        # Cursor-based pagination using (created_at, id) to avoid infinite loops
-        # and ensure proper ordering with time-based filtering
-        _cursor: tuple[datetime.datetime, str] | None = None
+        cursor: tuple[datetime.datetime, str] | None = None
+        app_to_tenant_cache: dict[str, str | None] = {}
+        current_candidate_batch_size = self._candidate_batch_size
+        smoothed_hit_rate: float | None = None
+        session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
 
         logger.info(
-            "clean_messages: start cleaning messages (dry_run=%s), start_from=%s, end_before=%s",
+            "clean_messages: start cleaning messages (dry_run=%s), start_from=%s, end_before=%s, "
+            "candidate_batch_size=%s, max_candidate_batch_size=%s, delete_batch_size=%s",
             self._dry_run,
             self._start_from,
             self._end_before,
+            self._candidate_batch_size,
+            self._max_candidate_batch_size,
+            self._delete_batch_size,
         )
 
-        max_batch_interval_ms = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL
-
         while True:
             stats["batches"] += 1
             batch_start = time.monotonic()
@@ -369,25 +436,25 @@ class MessagesCleanService:
             batch_deleted_messages = 0
 
             # Step 1: Fetch a batch of messages using cursor
-            with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
+            with session_factory.begin() as session:
                 fetch_messages_start = time.monotonic()
                 msg_stmt = (
                     select(Message.id, Message.app_id, Message.created_at)
                     .where(Message.created_at < self._end_before)
                     .order_by(Message.created_at, Message.id)
-                    .limit(self._batch_size)
+                    .limit(current_candidate_batch_size)
                 )
 
                 if self._start_from:
                     msg_stmt = msg_stmt.where(Message.created_at >= self._start_from)
 
                 # Apply cursor condition: (created_at, id) > (last_created_at, last_message_id)
-                if _cursor:
+                if cursor:
                     msg_stmt = msg_stmt.where(
                         tuple_(Message.created_at, Message.id)
                         > tuple_(
-                            sa.literal(_cursor[0], type_=sa.DateTime()),
-                            sa.literal(_cursor[1], type_=Message.id.type),
+                            sa.literal(cursor[0], type_=sa.DateTime()),
+                            sa.literal(cursor[1], type_=Message.id.type),
                         )
                     )
 
@@ -397,13 +464,13 @@ class MessagesCleanService:
                     for msg_id, app_id, msg_created_at in raw_messages
                 ]
                 logger.info(
-                    "clean_messages (batch %s): fetched %s messages in %sms",
+                    "clean_messages (batch %s): fetched %s candidate messages with limit %s in %sms",
                     stats["batches"],
                     len(messages),
+                    current_candidate_batch_size,
                     int((time.monotonic() - fetch_messages_start) * 1000),
                 )
 
-                # Track total messages fetched across all batches
                 stats["total_messages"] += len(messages)
                 batch_scanned_messages = len(messages)
 
@@ -417,29 +484,50 @@ class MessagesCleanService:
                     )
                     break
 
-                # Update cursor to the last message's (created_at, id)
-                _cursor = (messages[-1].created_at, messages[-1].id)
+                # Advance by candidate rows before policy filtering. This avoids retrying the same paid/unknown slice.
+                cursor = (messages[-1].created_at, messages[-1].id)
 
                 # Step 2: Extract app_ids and query tenant_ids
                 app_ids = list({msg.app_id for msg in messages})
 
                 if not app_ids:
                     logger.info("clean_messages (batch %s): no app_ids found, skip", stats["batches"])
+                    smoothed_hit_rate, current_candidate_batch_size = self._adjust_candidate_batch_size(
+                        smoothed_hit_rate=smoothed_hit_rate,
+                        candidate_count=batch_scanned_messages,
+                        eligible_count=0,
+                    )
+                    self._metrics.record_batch(
+                        scanned_messages=batch_scanned_messages,
+                        filtered_messages=batch_filtered_messages,
+                        deleted_messages=batch_deleted_messages,
+                        batch_duration_seconds=time.monotonic() - batch_start,
+                    )
                     continue
 
                 fetch_apps_start = time.monotonic()
-                app_stmt = select(App.id, App.tenant_id).where(App.id.in_(app_ids))
-                apps = list(session.execute(app_stmt).all())
+                app_to_tenant, app_cache_misses, apps_found = self._load_app_to_tenant_mapping(
+                    session=session,
+                    app_ids=app_ids,
+                    app_to_tenant_cache=app_to_tenant_cache,
+                )
                 logger.info(
-                    "clean_messages (batch %s): fetched %s apps for %s app_ids in %sms",
+                    "clean_messages (batch %s): resolved %s apps for %s app_ids (cache_misses=%s, found=%s) in %sms",
                     stats["batches"],
-                    len(apps),
+                    len(app_to_tenant),
                     len(app_ids),
+                    app_cache_misses,
+                    apps_found,
                     int((time.monotonic() - fetch_apps_start) * 1000),
                 )
 
-            if not apps:
+            if not app_to_tenant:
                 logger.info("clean_messages (batch %s): no apps found, skip", stats["batches"])
+                smoothed_hit_rate, current_candidate_batch_size = self._adjust_candidate_batch_size(
+                    smoothed_hit_rate=smoothed_hit_rate,
+                    candidate_count=batch_scanned_messages,
+                    eligible_count=0,
+                )
                 self._metrics.record_batch(
                     scanned_messages=batch_scanned_messages,
                     filtered_messages=batch_filtered_messages,
@@ -448,19 +536,27 @@ class MessagesCleanService:
                 )
                 continue
 
-            # Build app_id -> tenant_id mapping
-            app_to_tenant: dict[str, str] = {app.id: app.tenant_id for app in apps}
-
             # Step 3: Delegate to policy to determine which messages to delete
             policy_start = time.monotonic()
-            message_ids_to_delete = self._policy.filter_message_ids(messages, app_to_tenant)
+            message_ids_to_delete = list(self._policy.filter_message_ids(messages, app_to_tenant))
+            batch_filtered_messages = len(message_ids_to_delete)
+            stats["filtered_messages"] += batch_filtered_messages
+            smoothed_hit_rate, next_candidate_batch_size = self._adjust_candidate_batch_size(
+                smoothed_hit_rate=smoothed_hit_rate,
+                candidate_count=batch_scanned_messages,
+                eligible_count=batch_filtered_messages,
+            )
             logger.info(
-                "clean_messages (batch %s): policy selected %s/%s messages in %sms",
+                "clean_messages (batch %s): policy selected %s/%s messages in %sms "
+                "(smoothed_hit_rate=%.4f, next_candidate_batch_size=%s)",
                 stats["batches"],
-                len(message_ids_to_delete),
+                batch_filtered_messages,
                 len(messages),
                 int((time.monotonic() - policy_start) * 1000),
+                smoothed_hit_rate,
+                next_candidate_batch_size,
             )
+            current_candidate_batch_size = next_candidate_batch_size
 
             if not message_ids_to_delete:
                 logger.info("clean_messages (batch %s): no messages to delete, skip", stats["batches"])
@@ -472,47 +568,32 @@ class MessagesCleanService:
                 )
                 continue
 
-            stats["filtered_messages"] += len(message_ids_to_delete)
-            batch_filtered_messages = len(message_ids_to_delete)
-
             # Step 4: Batch delete messages and their relations
             if not self._dry_run:
-                with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
-                    delete_relations_start = time.monotonic()
-                    # Delete related records first
-                    self._batch_delete_message_relations(session, message_ids_to_delete)
-                    delete_relations_ms = int((time.monotonic() - delete_relations_start) * 1000)
+                delete_result = self._delete_messages_in_chunks(
+                    session_factory=session_factory,
+                    message_ids=message_ids_to_delete,
+                )
 
-                    # Delete messages
-                    delete_messages_start = time.monotonic()
-                    delete_stmt = delete(Message).where(Message.id.in_(message_ids_to_delete))
-                    delete_result = cast(CursorResult, session.execute(delete_stmt))
-                    messages_deleted = delete_result.rowcount
-                    delete_messages_ms = int((time.monotonic() - delete_messages_start) * 1000)
-                    commit_ms = 0
+                stats["total_deleted"] += delete_result["messages_deleted"]
+                batch_deleted_messages = delete_result["messages_deleted"]
 
-                    stats["total_deleted"] += messages_deleted
-                    batch_deleted_messages = messages_deleted
+                logger.info(
+                    "clean_messages (batch %s): processed %s candidate messages, deleted %s messages in %s chunks",
+                    stats["batches"],
+                    len(messages),
+                    delete_result["messages_deleted"],
+                    delete_result["chunks"],
+                )
+                logger.info(
+                    "clean_messages (batch %s): relations %sms, messages %sms, batch total %sms",
+                    stats["batches"],
+                    delete_result["relations_ms"],
+                    delete_result["messages_ms"],
+                    int((time.monotonic() - batch_start) * 1000),
+                )
 
-                    logger.info(
-                        "clean_messages (batch %s): processed %s messages, deleted %s messages",
-                        stats["batches"],
-                        len(messages),
-                        messages_deleted,
-                    )
-                    logger.info(
-                        "clean_messages (batch %s): relations %sms,  messages %sms, commit %sms, batch total %sms",
-                        stats["batches"],
-                        delete_relations_ms,
-                        delete_messages_ms,
-                        commit_ms,
-                        int((time.monotonic() - batch_start) * 1000),
-                    )
-
-                # Random sleep between batches to avoid overwhelming the database
-                sleep_ms = random.uniform(0, max_batch_interval_ms)  # noqa: S311
-                logger.info("clean_messages (batch %s): sleeping for %.2fms", stats["batches"], sleep_ms)
-                time.sleep(sleep_ms / 1000)
+                self._sleep_after_batch(stats["batches"], batch_deleted_messages)
             else:
                 # Log random sample of message IDs that would be deleted (up to 10)
                 sample_size = min(10, len(message_ids_to_delete))
@@ -544,6 +625,105 @@ class MessagesCleanService:
 
         return stats
 
+    @staticmethod
+    def _load_app_to_tenant_mapping(
+        *,
+        session: Session,
+        app_ids: Sequence[str],
+        app_to_tenant_cache: dict[str, str | None],
+    ) -> tuple[dict[str, str], int, int]:
+        unique_app_ids = sorted(set(app_ids))
+        missing_app_ids = [app_id for app_id in unique_app_ids if app_id not in app_to_tenant_cache]
+        found_apps = 0
+
+        if missing_app_ids:
+            app_stmt = select(App.id, App.tenant_id).where(App.id.in_(missing_app_ids))
+            apps = list(session.execute(app_stmt).all())
+            found_app_ids: set[str] = set()
+            for app_id, tenant_id in apps:
+                app_to_tenant_cache[app_id] = tenant_id
+                found_app_ids.add(app_id)
+            found_apps = len(found_app_ids)
+
+            for app_id in set(missing_app_ids) - found_app_ids:
+                app_to_tenant_cache[app_id] = None
+
+        app_to_tenant: dict[str, str] = {}
+        for app_id in unique_app_ids:
+            tenant_id = app_to_tenant_cache.get(app_id)
+            if tenant_id is not None:
+                app_to_tenant[app_id] = tenant_id
+
+        return app_to_tenant, len(missing_app_ids), found_apps
+
+    def _adjust_candidate_batch_size(
+        self,
+        *,
+        smoothed_hit_rate: float | None,
+        candidate_count: int,
+        eligible_count: int,
+    ) -> tuple[float, int]:
+        if candidate_count <= 0:
+            next_smoothed_hit_rate = smoothed_hit_rate or 0.0
+            return next_smoothed_hit_rate, self._candidate_batch_size_for_hit_rate(next_smoothed_hit_rate)
+
+        hit_rate = eligible_count / candidate_count
+        if smoothed_hit_rate is None:
+            next_smoothed_hit_rate = hit_rate
+        else:
+            next_smoothed_hit_rate = (_HIT_RATE_EMA_ALPHA * hit_rate) + ((1 - _HIT_RATE_EMA_ALPHA) * smoothed_hit_rate)
+
+        return next_smoothed_hit_rate, self._candidate_batch_size_for_hit_rate(next_smoothed_hit_rate)
+
+    def _candidate_batch_size_for_hit_rate(self, hit_rate: float) -> int:
+        effective_hit_rate = max(hit_rate, _MIN_ELIGIBLE_HIT_RATE)
+        desired_batch_size = math.ceil(self._delete_batch_size / effective_hit_rate)
+        lower_bound = min(self._delete_batch_size, self._max_candidate_batch_size)
+        return min(max(desired_batch_size, lower_bound), self._max_candidate_batch_size)
+
+    def _delete_messages_in_chunks(
+        self,
+        *,
+        session_factory: sessionmaker[Session],
+        message_ids: Sequence[str],
+    ) -> MessageDeleteResultDict:
+        result: MessageDeleteResultDict = {
+            "messages_deleted": 0,
+            "chunks": 0,
+            "relations_ms": 0,
+            "messages_ms": 0,
+        }
+
+        for message_id_chunk in self._iter_message_id_chunks(message_ids, self._delete_batch_size):
+            result["chunks"] += 1
+            with session_factory.begin() as session:
+                delete_relations_start = time.monotonic()
+                self._batch_delete_message_relations(session, message_id_chunk)
+                result["relations_ms"] += int((time.monotonic() - delete_relations_start) * 1000)
+
+                delete_messages_start = time.monotonic()
+                delete_stmt = delete(Message).where(Message.id.in_(message_id_chunk))
+                delete_result = cast(CursorResult, session.execute(delete_stmt))
+                result["messages_deleted"] += delete_result.rowcount
+                result["messages_ms"] += int((time.monotonic() - delete_messages_start) * 1000)
+
+        return result
+
+    @staticmethod
+    def _iter_message_id_chunks(message_ids: Sequence[str], chunk_size: int) -> Iterator[Sequence[str]]:
+        for start_index in range(0, len(message_ids), chunk_size):
+            yield message_ids[start_index : start_index + chunk_size]
+
+    def _sleep_after_batch(self, batch_index: int, deleted_messages: int) -> None:
+        if deleted_messages <= 0:
+            return
+
+        max_batch_interval_ms = dify_config.SANDBOX_EXPIRED_RECORDS_CLEAN_BATCH_MAX_INTERVAL
+        sleep_ratio = min(1.0, deleted_messages / self._delete_batch_size)
+        sleep_ms = random.uniform(0, max_batch_interval_ms * sleep_ratio)  # noqa: S311
+        logger.info("clean_messages (batch %s): sleeping for %.2fms", batch_index, sleep_ms)
+        time.sleep(sleep_ms / 1000)
+
     @staticmethod
     def _batch_delete_message_relations(session: Session, message_ids: Sequence[str]) -> None:
         """

api/tests/integration_tests/services/retention/test_messages_clean_service.py

@@ -1,11 +1,14 @@
 import datetime
 import math
 import uuid
+from unittest.mock import MagicMock, patch
 
 import pytest
 from sqlalchemy import delete, func, select
+from sqlalchemy.orm import Session
 
 from core.db.session_factory import session_factory
+from enums.cloud_plan import CloudPlan
 from models import Tenant
 from models.enums import FeedbackFromSource, FeedbackRating
 from models.model import (
@@ -15,7 +18,7 @@ from models.model import (
     MessageAnnotation,
     MessageFeedback,
 )
-from services.retention.conversation.messages_clean_policy import BillingDisabledPolicy
+from services.retention.conversation.messages_clean_policy import BillingDisabledPolicy, BillingSandboxPolicy
 from services.retention.conversation.messages_clean_service import MessagesCleanService
 
 _NOW = datetime.datetime(2026, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
@@ -91,6 +94,35 @@ def _make_message(app_id: str, conversation_id: str, created_at: datetime.dateti
     )
 
 
+def _create_tenant_app_conversation(session: Session, name_suffix: str) -> tuple[str, str, str]:
+    tenant = Tenant(name=f"retention_it_tenant_{name_suffix}")
+    session.add(tenant)
+    session.flush()
+
+    app = App(
+        tenant_id=tenant.id,
+        name=f"Retention IT App {name_suffix}",
+        mode="chat",
+        enable_site=True,
+        enable_api=True,
+    )
+    session.add(app)
+    session.flush()
+
+    conv = Conversation(
+        app_id=app.id,
+        mode="chat",
+        name=f"test_conv_{name_suffix}",
+        status="normal",
+        from_source="console",
+        _inputs={},
+    )
+    session.add(conv)
+    session.flush()
+
+    return tenant.id, app.id, conv.id
+
+
 class TestMessagesCleanServiceIntegration:
     @pytest.fixture
     def seed_messages(self, tenant_and_app):
@@ -285,6 +317,117 @@ class TestMessagesCleanServiceIntegration:
             remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids)))
         assert remaining == 0
 
+    def test_candidate_cursor_advances_when_first_batch_has_no_eligible_messages(self, flask_req_ctx):
+        """A paid-only candidate batch must not prevent later sandbox messages from being cleaned."""
+        del flask_req_ctx
+        with session_factory.create_session() as session:
+            paid_tenant_id, paid_app_id, paid_conv_id = _create_tenant_app_conversation(session, "paid")
+            free_tenant_id, free_app_id, free_conv_id = _create_tenant_app_conversation(session, "free")
+
+            paid_msg_1 = _make_message(paid_app_id, paid_conv_id, _OLD)
+            paid_msg_2 = _make_message(paid_app_id, paid_conv_id, _OLD + datetime.timedelta(seconds=1))
+            free_msg = _make_message(free_app_id, free_conv_id, _OLD + datetime.timedelta(seconds=2))
+            session.add_all([paid_msg_1, paid_msg_2, free_msg])
+            session.commit()
+
+            paid_message_ids = [paid_msg_1.id, paid_msg_2.id]
+            free_message_id = free_msg.id
+            app_ids = [paid_app_id, free_app_id]
+            conversation_ids = [paid_conv_id, free_conv_id]
+            tenant_ids = [paid_tenant_id, free_tenant_id]
+
+        plan_map = {
+            paid_tenant_id: {"plan": CloudPlan.PROFESSIONAL, "expiration_date": -1},
+            free_tenant_id: {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+        }
+        plan_provider = MagicMock(
+            side_effect=lambda tenant_ids: {tenant_id: plan_map[tenant_id] for tenant_id in tenant_ids}
+        )
+        policy = BillingSandboxPolicy(plan_provider=plan_provider, graceful_period_days=21)
+
+        try:
+            with patch("services.retention.conversation.messages_clean_service.time.sleep"):
+                svc = MessagesCleanService.from_time_range(
+                    policy=policy,
+                    start_from=_OLD - datetime.timedelta(seconds=1),
+                    end_before=_OLD + datetime.timedelta(seconds=3),
+                    batch_size=2,
+                    max_candidate_batch_size=2,
+                    delete_batch_size=1,
+                )
+                stats = svc.run()
+
+            assert stats["total_messages"] == 3
+            assert stats["filtered_messages"] == 1
+            assert stats["total_deleted"] == 1
+            assert plan_provider.call_count == 2
+            assert set(plan_provider.call_args_list[0].args[0]) == {paid_tenant_id}
+            assert set(plan_provider.call_args_list[1].args[0]) == {free_tenant_id}
+
+            with session_factory.create_session() as session:
+                remaining_paid = session.scalar(
+                    select(func.count()).select_from(Message).where(Message.id.in_(paid_message_ids))
+                )
+                remaining_free = session.scalar(
+                    select(func.count()).select_from(Message).where(Message.id == free_message_id)
+                )
+
+            assert remaining_paid == 2
+            assert remaining_free == 0
+        finally:
+            with session_factory.create_session() as session:
+                session.execute(delete(Message).where(Message.id.in_([*paid_message_ids, free_message_id])))
+                session.execute(delete(Conversation).where(Conversation.id.in_(conversation_ids)))
+                session.execute(delete(App).where(App.id.in_(app_ids)))
+                session.execute(delete(Tenant).where(Tenant.id.in_(tenant_ids)))
+                session.commit()
+
+    def test_delete_batch_size_chunks_eligible_message_deletes(self, tenant_and_app):
+        """Candidate scans can be larger than the delete transaction chunk size."""
+        data = tenant_and_app
+        app_id = data["app_id"]
+        conv_id = data["conversation_id"]
+        msg_ids: list[str] = []
+
+        with session_factory.create_session() as session:
+            for index in range(5):
+                msg = _make_message(app_id, conv_id, _OLD + datetime.timedelta(seconds=index))
+                session.add(msg)
+                session.flush()
+                msg_ids.append(msg.id)
+            session.commit()
+
+        try:
+            with (
+                patch.object(
+                    MessagesCleanService,
+                    "_batch_delete_message_relations",
+                    wraps=MessagesCleanService._batch_delete_message_relations,
+                ) as delete_relations,
+                patch("services.retention.conversation.messages_clean_service.time.sleep"),
+            ):
+                svc = MessagesCleanService.from_time_range(
+                    policy=BillingDisabledPolicy(),
+                    start_from=_OLD - datetime.timedelta(seconds=1),
+                    end_before=_OLD + datetime.timedelta(seconds=6),
+                    batch_size=5,
+                    max_candidate_batch_size=5,
+                    delete_batch_size=2,
+                )
+                stats = svc.run()
+
+            assert stats["total_deleted"] == 5
+            assert delete_relations.call_count == 3
+
+            with session_factory.create_session() as session:
+                remaining = session.scalar(select(func.count()).select_from(Message).where(Message.id.in_(msg_ids)))
+
+            assert remaining == 0
+        finally:
+            with session_factory.create_session() as session:
+                session.execute(delete(Message).where(Message.id.in_(msg_ids)))
+                session.commit()
+
     def test_no_messages_in_range_returns_empty_stats(self, seed_messages):
         """A window entirely in the future must yield zero matches."""
         far_future = _NOW + datetime.timedelta(days=365)

api/tests/unit_tests/services/test_messages_clean_service.py

@@ -296,6 +296,64 @@ class TestBillingSandboxPolicyFilterMessageIds:
         called_tenant_ids = set(plan_provider.call_args[0][0])
         assert called_tenant_ids == {"tenant1", "tenant2"}
 
+    def test_plan_provider_reuses_job_level_cache(self):
+        """Test that repeated tenants are not fetched from billing more than once."""
+        # Arrange
+        now = self.CURRENT_TIMESTAMP
+        tenant_plans = {
+            "tenant1": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+            "tenant2": {"plan": CloudPlan.SANDBOX, "expiration_date": -1},
+        }
+        plan_provider = MagicMock(
+            side_effect=lambda tenant_ids: {tenant_id: tenant_plans[tenant_id] for tenant_id in tenant_ids}
+        )
+        policy = BillingSandboxPolicy(
+            plan_provider=plan_provider,
+            graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
+            current_timestamp=now,
+        )
+
+        # Act
+        first_result = policy.filter_message_ids(
+            [make_simple_message("msg1", "app1")],
+            {"app1": "tenant1"},
+        )
+        second_result = policy.filter_message_ids(
+            [
+                make_simple_message("msg2", "app1"),
+                make_simple_message("msg3", "app2"),
+            ],
+            {"app1": "tenant1", "app2": "tenant2"},
+        )
+
+        # Assert
+        assert set(first_result) == {"msg1"}
+        assert set(second_result) == {"msg2", "msg3"}
+        assert plan_provider.call_count == 2
+        assert set(plan_provider.call_args_list[0].args[0]) == {"tenant1"}
+        assert set(plan_provider.call_args_list[1].args[0]) == {"tenant2"}
+
+    def test_whitelisted_tenants_are_not_fetched_from_plan_provider(self):
+        """Test that whitelisted tenants are skipped before billing plan lookup."""
+        # Arrange
+        plan_provider = make_plan_provider({})
+        policy = BillingSandboxPolicy(
+            plan_provider=plan_provider,
+            graceful_period_days=self.GRACEFUL_PERIOD_DAYS,
+            tenant_whitelist=["tenant1"],
+            current_timestamp=self.CURRENT_TIMESTAMP,
+        )
+
+        # Act
+        result = policy.filter_message_ids(
+            [make_simple_message("msg1", "app1")],
+            {"app1": "tenant1"},
+        )
+
+        # Assert
+        assert list(result) == []
+        plan_provider.assert_not_called()
+
     def test_complex_mixed_scenario(self):
         """Test complex scenario with mixed plans, expirations, whitelist, and missing mappings."""
         # Arrange
@@ -497,6 +555,22 @@ class TestMessagesCleanServiceFromTimeRange:
                 batch_size=-100,
             )
 
+        with pytest.raises(ValueError, match="max_candidate_batch_size .* must be greater than 0"):
+            MessagesCleanService.from_time_range(
+                policy=policy,
+                start_from=start_from,
+                end_before=end_before,
+                max_candidate_batch_size=0,
+            )
+
+        with pytest.raises(ValueError, match="delete_batch_size .* must be greater than 0"):
+            MessagesCleanService.from_time_range(
+                policy=policy,
+                start_from=start_from,
+                end_before=end_before,
+                delete_batch_size=0,
+            )
+
     def test_valid_params_creates_instance(self):
         """Test that valid parameters create a correctly configured instance."""
         # Arrange
@@ -512,6 +586,8 @@ class TestMessagesCleanServiceFromTimeRange:
             start_from=start_from,
             end_before=end_before,
             batch_size=batch_size,
+            max_candidate_batch_size=5000,
+            delete_batch_size=200,
             dry_run=dry_run,
         )
 
@@ -521,6 +597,9 @@ class TestMessagesCleanServiceFromTimeRange:
         assert service._start_from == start_from
         assert service._end_before == end_before
         assert service._batch_size == batch_size
+        assert service._candidate_batch_size == batch_size
+        assert service._max_candidate_batch_size == 5000
+        assert service._delete_batch_size == 200
         assert service._dry_run == dry_run
 
     def test_default_params(self):
@@ -539,6 +618,9 @@ class TestMessagesCleanServiceFromTimeRange:
 
         # Assert
         assert service._batch_size == 1000  # default
+        assert service._candidate_batch_size == 1000  # default
+        assert service._max_candidate_batch_size == 1000  # default
+        assert service._delete_batch_size == 1000  # default
         assert service._dry_run is False  # default
 
     def test_explicit_task_label(self):
@@ -590,6 +672,12 @@ class TestMessagesCleanServiceFromDays:
         with pytest.raises(ValueError, match="batch_size .* must be greater than 0"):
             MessagesCleanService.from_days(policy=policy, days=30, batch_size=-500)
 
+        with pytest.raises(ValueError, match="max_candidate_batch_size .* must be greater than 0"):
+            MessagesCleanService.from_days(policy=policy, days=30, max_candidate_batch_size=0)
+
+        with pytest.raises(ValueError, match="delete_batch_size .* must be greater than 0"):
+            MessagesCleanService.from_days(policy=policy, days=30, delete_batch_size=0)
+
     def test_valid_params_creates_instance(self):
         """Test that valid parameters create a correctly configured instance."""
         # Arrange
@@ -606,6 +694,8 @@ class TestMessagesCleanServiceFromDays:
                 policy=policy,
                 days=days,
                 batch_size=batch_size,
+                max_candidate_batch_size=2500,
+                delete_batch_size=250,
                 dry_run=dry_run,
             )
 
@@ -616,6 +706,9 @@ class TestMessagesCleanServiceFromDays:
         assert service._start_from is None
         assert service._end_before == expected_end_before
         assert service._batch_size == batch_size
+        assert service._candidate_batch_size == batch_size
+        assert service._max_candidate_batch_size == 2500
+        assert service._delete_batch_size == 250
         assert service._dry_run == dry_run
 
     def test_default_params(self):
@@ -637,6 +730,92 @@ class TestMessagesCleanServiceFromDays:
         assert service._metrics._base_attributes["task_label"] == "custom"
 
 
+class TestMessagesCleanServiceBatchHelpers:
+    """Unit tests for cache and adaptive batch helpers."""
+
+    def test_load_app_to_tenant_mapping_reuses_cache(self):
+        class ExecuteResult:
+            def __init__(self, rows: list[tuple[str, str]]) -> None:
+                self._rows = rows
+
+            def all(self) -> list[tuple[str, str]]:
+                return self._rows
+
+        class FakeSession:
+            execute_calls: int
+
+            def __init__(self) -> None:
+                self.execute_calls = 0
+
+            def execute(self, _stmt: object) -> ExecuteResult:
+                self.execute_calls += 1
+                return ExecuteResult([("app1", "tenant1")])
+
+        session = FakeSession()
+        cache: dict[str, str | None] = {}
+
+        first_mapping, first_cache_misses, first_found = MessagesCleanService._load_app_to_tenant_mapping(
+            session=session,  # type: ignore[arg-type]
+            app_ids=["app1"],
+            app_to_tenant_cache=cache,
+        )
+        second_mapping, second_cache_misses, second_found = MessagesCleanService._load_app_to_tenant_mapping(
+            session=session,  # type: ignore[arg-type]
+            app_ids=["app1"],
+            app_to_tenant_cache=cache,
+        )
+
+        assert first_mapping == {"app1": "tenant1"}
+        assert second_mapping == {"app1": "tenant1"}
+        assert first_cache_misses == 1
+        assert first_found == 1
+        assert second_cache_misses == 0
+        assert second_found == 0
+        assert session.execute_calls == 1
+
+    def test_candidate_batch_size_grows_for_low_hit_rate(self):
+        service = MessagesCleanService(
+            policy=BillingDisabledPolicy(),
+            start_from=datetime.datetime(2024, 1, 1),
+            end_before=datetime.datetime(2024, 1, 2),
+            batch_size=1000,
+            max_candidate_batch_size=50000,
+            delete_batch_size=1000,
+        )
+
+        smoothed_hit_rate, next_batch_size = service._adjust_candidate_batch_size(
+            smoothed_hit_rate=None,
+            candidate_count=10000,
+            eligible_count=55,
+        )
+
+        assert smoothed_hit_rate == 0.0055
+        assert next_batch_size == 50000
+
+    def test_candidate_batch_size_shrinks_when_hit_rate_is_high(self):
+        service = MessagesCleanService(
+            policy=BillingDisabledPolicy(),
+            start_from=datetime.datetime(2024, 1, 1),
+            end_before=datetime.datetime(2024, 1, 2),
+            batch_size=10000,
+            max_candidate_batch_size=50000,
+            delete_batch_size=1000,
+        )
+
+        _smoothed_hit_rate, next_batch_size = service._adjust_candidate_batch_size(
+            smoothed_hit_rate=None,
+            candidate_count=10000,
+            eligible_count=10000,
+        )
+
+        assert next_batch_size == 1000
+
+    def test_iter_message_id_chunks_uses_delete_batch_size(self):
+        chunks = list(MessagesCleanService._iter_message_id_chunks(["msg1", "msg2", "msg3", "msg4", "msg5"], 2))
+
+        assert chunks == [["msg1", "msg2"], ["msg3", "msg4"], ["msg5"]]
+
+
 class TestMessagesCleanServiceRun:
     """Unit tests for MessagesCleanService.run instrumentation behavior."""