resolutionflow/backend/app/services/retention_cleanup.py

"""Chat retention cleanup job.

Runs daily via APScheduler to enforce account-level retention settings:
- Delete non-pinned chats older than chat_retention_days
- Delete oldest non-pinned chats when count exceeds chat_retention_max_count
"""
import logging
from datetime import datetime, timezone, timedelta

from sqlalchemy import select, delete, func

from app.core.database import async_session_maker
from app.models.account import Account
from app.models.assistant_chat import AssistantChat

logger = logging.getLogger(__name__)


async def cleanup_expired_chats() -> None:
    """Enforce chat retention policies for all accounts."""
    async with async_session_maker() as db:
        try:
            result = await db.execute(select(Account))
            accounts = result.scalars().all()

            total_deleted = 0
            for account in accounts:
                deleted = await _cleanup_account_chats(account, db)
                total_deleted += deleted

            await db.commit()
            if total_deleted > 0:
                logger.info("[retention] Cleaned up %d expired chats", total_deleted)
        except Exception as e:
            logger.error("[retention] Chat cleanup failed: %s", e)
            await db.rollback()


async def _cleanup_account_chats(account: Account, db) -> int:
    """Enforce retention for a single account. Returns count deleted."""
    deleted = 0

    # Age-based retention
    if account.chat_retention_days:
        cutoff = datetime.now(timezone.utc) - timedelta(days=account.chat_retention_days)
        result = await db.execute(
            delete(AssistantChat)
            .where(
                AssistantChat.account_id == account.id,
                AssistantChat.pinned == False,  # noqa: E712
                AssistantChat.updated_at < cutoff,
            )
            .returning(AssistantChat.id)
        )
        deleted += len(result.all())

    # Count-based retention
    if account.chat_retention_max_count:
        total = await db.scalar(
            select(func.count(AssistantChat.id)).where(
                AssistantChat.account_id == account.id,
            )
        ) or 0

        if total > account.chat_retention_max_count:
            excess = total - account.chat_retention_max_count
            # Get oldest non-pinned chat IDs
            oldest = await db.execute(
                select(AssistantChat.id)
                .where(
                    AssistantChat.account_id == account.id,
                    AssistantChat.pinned == False,  # noqa: E712
                )
                .order_by(AssistantChat.updated_at.asc())
                .limit(excess)
            )
            ids_to_delete = [row[0] for row in oldest.all()]
            if ids_to_delete:
                await db.execute(
                    delete(AssistantChat).where(AssistantChat.id.in_(ids_to_delete))
                )
                deleted += len(ids_to_delete)

    return deleted