"""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)