odysseus/src/cleanup_service.py

# src/cleanup_service.py
import logging
from datetime import datetime, timedelta, timezone
from typing import Tuple, Dict, Any, Optional

logger = logging.getLogger(__name__)


def _utcnow() -> datetime:
    """Naive UTC for this module's DB-bound timestamps.

    Mirrors the naive DateTime columns these values are compared against,
    without the deprecated stdlib UTC-now call (removed in Python 3.14).
    """
    return datetime.now(timezone.utc).replace(tzinfo=None)


class CleanupConfig:
    """Configuration constants for cleanup operations."""
    ARCHIVE_AFTER_DAYS = 7
    DELETE_AFTER_DAYS = 14
    MIN_MESSAGES_TO_KEEP = 20
    PRESERVE_RECENT_COUNT = 10
    PROTECTED_KEYWORDS = ['important', 'remember', 'save this', 'keep', 'bookmark']
    ESTIMATED_MESSAGE_SIZE_BYTES = 512


def _apply_owner_filter(query, DbSession, owner: Optional[str]):
    """Apply owner filtering to a session query.

    SECURITY: strict — the previous OR predicate let one user's cleanup
    archive/delete every null-owner session, including ones that hadn't
    been migrated. Now: only rows owned by this user.
    """
    if owner is None:
        return query
    return query.filter(DbSession.owner == owner)


async def archive_inactive_sessions(session_manager, owner: Optional[str] = None) -> int:
    """
    Archive sessions that haven't been accessed in the configured number of days.

    Args:
        session_manager: The session manager instance
        owner: If set, only archive this user's sessions

    Returns:
        Number of sessions archived
    """
    cutoff_date = _utcnow() - timedelta(days=CleanupConfig.ARCHIVE_AFTER_DAYS)
    archived_count = 0

    from src.database import SessionLocal, Session as DbSession
    db = SessionLocal()
    try:
        q = db.query(DbSession).filter(
            DbSession.last_accessed < cutoff_date,
            DbSession.archived == False
        )
        q = _apply_owner_filter(q, DbSession, owner)
        sessions_to_archive = q.all()

        for session in sessions_to_archive:
            session.archived = True
            session.updated_at = _utcnow()
            archived_count += 1

        if archived_count > 0:
            db.commit()
            logger.info(f"Archived {archived_count} inactive sessions")

    except Exception as e:
        logger.error(f"Error archiving sessions: {e}")
        db.rollback()
    finally:
        db.close()

    return archived_count

async def cleanup_old_sessions(session_manager, owner: Optional[str] = None) -> Tuple[int, float]:
    """
    Delete old sessions based on specific criteria.

    Args:
        session_manager: The session manager instance
        owner: If set, only clean up this user's sessions

    Returns:
        Tuple of (number of sessions deleted, space freed in MB)
    """
    cutoff_date = _utcnow() - timedelta(days=CleanupConfig.DELETE_AFTER_DAYS)
    deleted_count = 0
    space_freed = 0

    from src.database import SessionLocal, Session as DbSession, ChatMessage as DbChatMessage
    db = SessionLocal()
    try:
        recent_q = db.query(DbSession).order_by(DbSession.created_at.desc())
        recent_q = _apply_owner_filter(recent_q, DbSession, owner)
        all_sessions = recent_q.all()
        recent_session_ids = {session.id for session in all_sessions[:CleanupConfig.PRESERVE_RECENT_COUNT]}

        base_query = db.query(DbSession).filter(
            DbSession.archived == True,
            DbSession.last_accessed < cutoff_date,
            DbSession.is_important == False,
            DbSession.message_count < CleanupConfig.MIN_MESSAGES_TO_KEEP
        )
        base_query = _apply_owner_filter(base_query, DbSession, owner)

        candidate_sessions = base_query.all()
        sessions_to_delete = []
        preserved_count = 0

        for session in candidate_sessions:
            if session.id in recent_session_ids:
                preserved_count += 1
                continue

            if session.message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP:
                preserved_count += 1
                continue

            session_name_lower = session.name.lower() if session.name else ""
            if any(keyword in session_name_lower for keyword in CleanupConfig.PROTECTED_KEYWORDS):
                preserved_count += 1
                continue

            sessions_to_delete.append(session)

        for session in sessions_to_delete:
            message_count = db.query(DbChatMessage).filter(
                DbChatMessage.session_id == session.id
            ).count()
            space_freed += message_count * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES

        session_ids = [session.id for session in sessions_to_delete]
        if session_ids:
            db.query(DbSession).filter(DbSession.id.in_(session_ids)).delete(synchronize_session=False)
            deleted_count = len(session_ids)
            db.commit()

            for session_id in session_ids:
                if session_id in session_manager.sessions:
                    del session_manager.sessions[session_id]

        if deleted_count > 0:
            space_freed_mb = space_freed / (1024 * 1024)
            logger.info(f"Deleted {deleted_count} old sessions, freeing approximately {space_freed_mb:.2f} MB")
            return deleted_count, space_freed_mb

    except Exception as e:
        logger.error(f"Error cleaning up old sessions: {e}")
        db.rollback()
    finally:
        db.close()

    return deleted_count, 0.0

async def get_cleanup_preview(owner: Optional[str] = None) -> Dict[str, Any]:
    """
    Get a preview of what would be cleaned up without making changes.

    Args:
        owner: If set, only preview this user's sessions

    Returns:
        Dictionary containing preview information
    """
    cutoff_archive = _utcnow() - timedelta(days=CleanupConfig.ARCHIVE_AFTER_DAYS)
    cutoff_delete = _utcnow() - timedelta(days=CleanupConfig.DELETE_AFTER_DAYS)

    sessions_to_archive = []
    sessions_to_delete = []
    estimated_space_freed = 0
    preserved_sessions = []

    from src.database import SessionLocal, Session as DbSession
    db = SessionLocal()
    try:
        archive_q = db.query(DbSession).filter(
            DbSession.last_accessed < cutoff_archive,
            DbSession.archived == False
        )
        archive_q = _apply_owner_filter(archive_q, DbSession, owner)
        archive_candidates = archive_q.all()

        for session in archive_candidates:
            sessions_to_archive.append({
                "id": session.id,
                "name": session.name,
                "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown",
                "message_count": session.message_count
            })

        recent_q = db.query(DbSession).order_by(DbSession.created_at.desc())
        recent_q = _apply_owner_filter(recent_q, DbSession, owner)
        all_sessions = recent_q.all()
        recent_session_ids = {session.id for session in all_sessions[:CleanupConfig.PRESERVE_RECENT_COUNT]}

        base_query = db.query(DbSession).filter(
            DbSession.archived == True,
            DbSession.last_accessed < cutoff_delete,
            DbSession.is_important == False,
            DbSession.message_count < CleanupConfig.MIN_MESSAGES_TO_KEEP
        )
        base_query = _apply_owner_filter(base_query, DbSession, owner)

        candidate_sessions = base_query.all()

        for session in candidate_sessions:
            if session.id in recent_session_ids:
                preserved_sessions.append({
                    "id": session.id,
                    "name": session.name,
                    "reason": f"part of last {CleanupConfig.PRESERVE_RECENT_COUNT} sessions",
                    "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown",
                    "message_count": session.message_count
                })
                continue

            if session.message_count >= CleanupConfig.MIN_MESSAGES_TO_KEEP:
                preserved_sessions.append({
                    "id": session.id,
                    "name": session.name,
                    "reason": f"has {CleanupConfig.MIN_MESSAGES_TO_KEEP}+ messages",
                    "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown",
                    "message_count": session.message_count
                })
                continue

            session_name_lower = session.name.lower() if session.name else ""
            matching_keywords = [keyword for keyword in CleanupConfig.PROTECTED_KEYWORDS if keyword in session_name_lower]
            if matching_keywords:
                preserved_sessions.append({
                    "id": session.id,
                    "name": session.name,
                    "reason": f"contains keyword: {matching_keywords[0]}",
                    "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown",
                    "message_count": session.message_count
                })
                continue

            session_space = session.message_count * CleanupConfig.ESTIMATED_MESSAGE_SIZE_BYTES
            estimated_space_freed += session_space

            sessions_to_delete.append({
                "id": session.id,
                "name": session.name,
                "last_accessed": session.last_accessed.isoformat() if session.last_accessed else "Unknown",
                "message_count": session.message_count,
                "estimated_size_kb": round(session_space / 1024, 2)
            })

    except Exception as e:
        logger.error(f"Error generating cleanup preview: {e}")
    finally:
        db.close()

    return {
        "sessions_to_archive": sessions_to_archive,
        "sessions_to_delete": sessions_to_delete,
        "preserved_sessions": preserved_sessions,
        "estimated_space_freed_mb": round(estimated_space_freed / (1024 * 1024), 2)
    }

async def cleanup_sessions(session_manager, owner: Optional[str] = None) -> Tuple[int, int, float]:
    """
    Perform complete cleanup operations with error recovery.

    Args:
        session_manager: The session manager instance
        owner: If set, only clean up this user's sessions

    Returns:
        Tuple of (archived_count, deleted_count, space_freed_mb)
    """
    archived_count = 0
    deleted_count = 0
    space_freed_mb = 0.0

    try:
        archived_count = await archive_inactive_sessions(session_manager, owner=owner)
    except Exception as e:
        logger.error(f"Archive operation failed: {e}")

    try:
        deleted_count, space_freed_mb = await cleanup_old_sessions(session_manager, owner=owner)
    except Exception as e:
        logger.error(f"Delete operation failed: {e}")

    return archived_count, deleted_count, space_freed_mb