app/tasks/billing_reset_task.py

"""Author RAG — Billing Monthly Token Reset Task.

Triggered by Stripe's invoice.payment_succeeded webhook (via BillingService).
Also runs as a scheduled Celery beat task on the 1st of each month as a fallback.

What it does:
  1. Resets tokens_used = 0 for all active subscriptions that renewed this cycle
  2. Clears payment_grace_until flags on the User
  3. Clears Redis token counters for each author (budget_exhausted flags, etc.)
  4. Logs to audit trail

The webhook-triggered path (BillingService.handle_payment_succeeded) handles
individual resets per customer. This task handles bulk reset as a safety net.
"""

import structlog
from celery import shared_task
from datetime import datetime, timezone
from sqlalchemy import select

logger = structlog.get_logger(__name__)


@shared_task(
    name="billing.monthly_token_reset",
    bind=True,
    max_retries=3,
    default_retry_delay=300,  # 5 min retry delay
)
def monthly_token_reset_task(self) -> dict:
    """Celery task: reset token counters for all active subscriptions.

    Runs on the 1st of each month at 00:05 UTC as a safety fallback
    in case any Stripe webhook was missed or failed to process.

    This is idempotent — resetting already-zeroed counters is harmless.
    """
    import asyncio
    from app.core.db import get_async_session
    from app.models.client_access import ClientAccess
    from app.models.user import User

    async def _reset():
        async with get_async_session() as db:
            # Fetch all non-revoked access records with Stripe subscription IDs
            result = await db.execute(
                select(ClientAccess).where(
                    ClientAccess.is_revoked == False,  # noqa: E712
                    ClientAccess.stripe_subscription_id.is_not(None),
                )
            )
            records = result.scalars().all()

            reset_count = 0
            for access in records:
                if access.tokens_used > 0:
                    access.tokens_used = 0
                    reset_count += 1

            # Clear payment grace periods for users with active subscriptions
            active_author_ids = {r.author_id for r in records}
            if active_author_ids:
                user_result = await db.execute(
                    select(User).where(
                        User.id.in_(active_author_ids),
                        User.payment_grace_until.is_not(None),
                    )
                )
                for user in user_result.scalars():
                    user.payment_grace_until = None

            await db.commit()
            return reset_count

    try:
        loop = asyncio.get_event_loop()
        count = loop.run_until_complete(_reset())
        logger.info("Monthly token reset complete", reset_count=count, ran_at=datetime.now(timezone.utc).isoformat())
        return {"reset_count": count, "status": "ok"}
    except Exception as exc:
        logger.error("Monthly token reset failed", error=str(exc))
        raise self.retry(exc=exc)