smart-org/apps/api/app/workers/billing.py

"""Billing worker — rolls the usage ledger into draft invoices.

The usage ledger is the source of truth for billing; an invoice is a *snapshot*
of one closed period. That ordering matters: the worker never invents an amount
the ledger does not already contain, and it never talks to a payment provider —
moving money is the API's job (``BillingService.settle_payment``), not a
background job's.

The one hard requirement is **idempotency**. A cron job that runs twice for the
same period must not bill twice, so a period that already has an invoice is
skipped rather than re-created.

The same pass also expires memory lessons past their date (``expire_memories``).
That is housekeeping, not billing, but this is the one scheduled job every
deployment runs, and expiry needs no model provider.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta

from app.domain.base import utcnow
from app.domain.billing.models import Invoice, InvoiceStatus
from app.infrastructure.repositories.audit import UsageLedgerRepository
from app.infrastructure.repositories.billing import InvoiceRepository
from app.infrastructure.repositories.identity import OrganizationRepository
from app.infrastructure.repositories.memory import MemoryRepository
from app.workers.base import WorkerOutcome, worker_session

#: A period shorter than this is almost certainly a bug (a stray manual run),
#: not a real billing cycle. Refusing it keeps the invoice table meaningful.
MIN_PERIOD = timedelta(hours=1)


class BillingWorker:
    """Turns usage into invoices, once per period, per organization."""

    name = "billing-worker"

    async def roll_up(
        self,
        *,
        period_start: datetime | None = None,
        period_end: datetime | None = None,
        organization_id: str | None = None,
    ) -> WorkerOutcome:
        """Create a DRAFT invoice for every organization with usage in the period.

        ``organization_id`` narrows the pass to one tenant (useful for a
        backfill or a support-driven re-run).
        """
        end = period_end or utcnow()
        start = period_start or (end - timedelta(days=30))

        if end <= start:
            raise ValueError("period_end must be after period_start")
        if end - start < MIN_PERIOD:
            raise ValueError(
                f"billing period must be at least {MIN_PERIOD} long"
            )

        outcome = WorkerOutcome(name=self.name)

        async with worker_session() as session:
            invoices = InvoiceRepository(session)
            ledger = UsageLedgerRepository(session)
            organizations = OrganizationRepository(session)

            if organization_id is not None:
                targets = [organization_id]
            else:
                targets = [org.id for org in await organizations.list_all()]

            for org_id in targets:
                existing = await invoices.list_for_org(org_id)
                if any(self._covers(i, start, end) for i in existing):
                    outcome.skipped += 1
                    continue

                breakdown = await ledger.breakdown(
                    organization_id=org_id, since=start, until=end
                )
                # Six decimals, not two: a single small model call costs a
                # fraction of a cent, and rounding to cents here would silently
                # turn real usage into a zero invoice.
                amount = round(breakdown.total_price, 6)
                if amount <= 0:
                    # No billable usage: do not litter the table with $0 drafts.
                    outcome.skipped += 1
                    outcome.note(f"{org_id}: no billable usage in this period")
                    continue

                invoice = await invoices.create(
                    Invoice(
                        organization_id=org_id,
                        amount=amount,
                        status=InvoiceStatus.DRAFT,
                        period_start=start,
                        period_end=end,
                    )
                )
                outcome.processed += 1
                outcome.note(f"{org_id}: {invoice.id} amount={amount}")

        return outcome

    async def expire_memories(self) -> WorkerOutcome:
        """Soft-delete organization lessons whose expiry has passed.

        Pinned entries never expire, and rules and preferences carry no expiry at
        all, so only disposable lessons are touched. The sweep covers every
        organization at once, which is what a platform-wide schedule wants.
        """
        outcome = WorkerOutcome(name=self.name)
        async with worker_session() as session:
            expired = await MemoryRepository(session).expire(None)
        outcome.processed = expired
        outcome.note(f"memory: {expired} expired lesson(s) soft-deleted")
        return outcome

    @staticmethod
    def _naive_utc(value: datetime) -> datetime:
        """Normalise a timestamp to naive UTC for comparison.

        SQLite drops ``tzinfo`` on round-trip while Postgres keeps it, so
        comparing a freshly computed aware timestamp against one read back from
        the database raises ``TypeError`` on one backend and works on the other.
        Comparing in a single explicit frame avoids depending on which database
        the deployment happens to use.
        """
        if value.tzinfo is None:
            return value
        return value.astimezone(UTC).replace(tzinfo=None)

    @classmethod
    def _covers(cls, invoice: Invoice, start: datetime, end: datetime) -> bool:
        """Is this period already invoiced?

        Any invoice whose period *overlaps* the requested one counts. Overlap
        rather than exact equality, because two runs a few seconds apart would
        otherwise produce two near-identical invoices for the same window.
        """
        if invoice.period_start is None or invoice.period_end is None:
            return False
        window_start, window_end = cls._naive_utc(start), cls._naive_utc(end)
        return (
            cls._naive_utc(invoice.period_start) < window_end
            and window_start < cls._naive_utc(invoice.period_end)
        )


__all__ = ["MIN_PERIOD", "BillingWorker"]