life-echo/api/app/tasks/memory_compaction_tasks.py

"""Celery：memory compaction（近重复 chunk 软排除）。"""

from __future__ import annotations

import time
from datetime import datetime
from typing import Any

from celery import shared_task

from app.core.config import settings
from app.core.db import get_sync_db
from app.core.logging import get_logger
from app.core.memory_compaction_schedule import (
    finalize_memory_compaction_run,
    read_debounce_deadline_ts,
    release_scheduler_gate,
    schedule_memory_compaction_run,
    set_incremental_cursor_pair,
)
from app.core.redis_lock import acquire_redis_lock, release_redis_lock
from app.features.memory.compaction_service import run_memory_compaction_sync
from app.features.memory.repo import list_users_with_recent_chunks_sync

logger = get_logger(__name__)


@shared_task
def memory_compaction_sweep() -> dict[str, Any]:
    """Beat：为近期有记忆写入的用户调度 compaction（debounce 仍由 schedule 合并）。"""
    if not settings.memory_compaction_enabled:
        return {"skipped": True, "reason": "disabled"}
    hours = int(settings.memory_compaction_sweep_recent_hours)
    with get_sync_db() as session:
        user_ids = list_users_with_recent_chunks_sync(session, hours=hours)
    ctx_base: dict[str, Any] = {"trigger_source": "beat", "sweep_hours": hours}
    for uid in user_ids:
        schedule_memory_compaction_run(uid, dict(ctx_base))
    logger.info(
        "memory_compaction_sweep hours={} scheduled_users={}", hours, len(user_ids)
    )
    return {"scheduled": len(user_ids), "user_ids": user_ids}


@shared_task(bind=True, max_retries=12, default_retry_delay=20)
def memory_compaction_run(
    self, user_id: str, context: dict[str, Any] | None = None
) -> dict[str, Any]:
    if not settings.memory_compaction_enabled:
        return {"skipped": True, "reason": "disabled"}

    ctx = dict(context or {})
    deadline = read_debounce_deadline_ts(user_id)
    now = time.time()
    if deadline is not None and now < deadline:
        delay = max(1.0, deadline - now)
        raise self.retry(countdown=int(delay))

    lock = acquire_redis_lock(
        f"lock:memory_compaction:{user_id}",
        ttl_seconds=settings.memory_compaction_lock_ttl_seconds,
    )
    if lock is None:
        logger.info(
            "memory_compaction_skipped user_id={} skipped_reason=lock_not_acquired",
            user_id,
        )
        out = {"skipped": True, "reason": "lock_not_acquired"}
        finalize_memory_compaction_run(
            user_id,
            observed_deadline_ts=deadline,
            context=ctx,
        )
        return out

    try:
        with get_sync_db() as session:
            out = run_memory_compaction_sync(session, user_id, ctx)
            session.commit()

        if out.get("new_cursor_ts") and out.get("new_cursor_id") is not None:
            set_incremental_cursor_pair(
                user_id,
                datetime.fromisoformat(out["new_cursor_ts"]),
                str(out["new_cursor_id"]),
            )
        finalize_memory_compaction_run(
            user_id,
            observed_deadline_ts=deadline,
            context=ctx,
        )
        return out
    except Exception as exc:
        logger.warning("memory_compaction_run failed user_id={} err={}", user_id, exc)
        release_scheduler_gate(user_id)
        raise self.retry(exc=exc)
    finally:
        release_redis_lock(lock)