plm-lottery/app/wallet/balance.py

from embit.transaction import Transaction
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models import PendingTransaction, User, UtxoEvent
from app.wallet.plm_network import PLM_MAINNET


async def recompute_balance(session: AsyncSession, user_id: int) -> int:
    """Source of truth: sum of this user's confirmed, unspent UTXOs. Updates and
    returns the read-cache column (User.cached_balance_sats). Must be called
    within the same transaction as whatever inserted/updated utxo_events rows."""
    balance = await session.scalar(
        select(func.sum(UtxoEvent.amount_sats)).where(UtxoEvent.user_id == user_id, UtxoEvent.spent_txid.is_(None))
    )
    user = await session.get(User, user_id)
    user.cached_balance_sats = balance or 0
    return user.cached_balance_sats


async def compute_pending_balance(session: AsyncSession, user: User) -> tuple[int, bool]:
    """Balance including the user's own change still in flight.

    Placing a bet or a withdrawal spends whatever UTXOs cover the amount — often
    much larger than the amount actually moving, since select_utxos() picks
    whole UTXOs — and recompute_balance() drops that entire input total from
    cached_balance_sats the moment the tx is broadcast (spent_txid is set right
    away, well before the tx has any confirmations). The change output that
    returns the difference only becomes a UtxoEvent (and so counts toward the
    confirmed balance again) once it gets its own 1st confirmation. In between,
    User.cached_balance_sats alone understates the user's real balance by the
    full unconfirmed change amount, which can look like a much bigger loss than
    the tx actually represents.

    This walks every in-flight (status="pending") bet/withdrawal PendingTransaction
    of this user, decodes its current raw tx (kept up to date across RBF bumps —
    see tx/broadcast.py:bump_fee), and sums whichever outputs pay back to the
    user's own address. Adding that to cached_balance_sats gives the balance the
    user will end up with once everything currently in flight confirms.

    Returns (pending_inclusive_balance_sats, has_pending) — has_pending tells the
    caller whether this differs from the confirmed-only balance at all.
    """
    pending = (
        await session.scalars(
            select(PendingTransaction).where(
                PendingTransaction.user_id == user.id,
                PendingTransaction.kind.in_(("bet", "withdrawal")),
                # "building" as well as "pending": a building row's UTXOs are already
                # marked spent (see place_bet's two phases), so leaving it out would
                # make the displayed balance dip for the duration of the broadcast.
                PendingTransaction.status.in_(("building", "pending")),
            )
        )
    ).all()

    pending_change_sats = 0
    for row in pending:
        tx = Transaction.parse(bytes.fromhex(row.raw_tx_hex))
        for out in tx.vout:
            if out.script_pubkey.address(network=PLM_MAINNET) == user.address:
                pending_change_sats += out.value

    return user.cached_balance_sats + pending_change_sats, bool(pending)