"""Quota arithmetic — pure logic, no IO (架构文档 §7.4 V1 配额三段式). The three-phase protocol lives here as data transformations so it can be unit tested without a database: reserve → 冻结额度并写 bl_ledger(status=reserved) commit → 确认扣除 (status=committed) refund → 释放额度并写反向流水 (tx_type=refund) V1 doesn't involve money, but 用户对「次数被白扣」的敏感度不比扣钱低,and V1 has no customer-service refund path — the mechanism is the only guarantee. """ from __future__ import annotations from dataclasses import dataclass, replace from datetime import UTC, datetime, timedelta, timezone from typing import Literal, get_args from packages.domain.errors import QuotaExceededError PeriodType = Literal["day", "month", "lifetime"] #: 合法周期取值集合,供配置校验复用;单一事实来源是上面的 Literal。 PERIOD_TYPES: frozenset[str] = frozenset(get_args(PeriodType)) LedgerStatus = Literal["reserved", "committed", "refunded"] TxType = Literal["system_grant", "consume", "refund", "purchase", "gift"] #: FR-BILL-01: 跨时区按 UTC+8 统一。 DEFAULT_TZ_OFFSET_HOURS = 8 def period_key( now: datetime, period_type: PeriodType, tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS, ) -> str: """Bucket key that makes quota reset naturally, with no cron job. ``day`` → ``d:2026-08-13``; ``month`` → ``m:2026-08``; ``lifetime`` → ``life``. """ if period_type == "lifetime": return "life" local = now.astimezone(UTC) + timedelta(hours=tz_offset_hours) if period_type == "day": return f"d:{local:%Y-%m-%d}" if period_type == "month": return f"m:{local:%Y-%m}" msg = f"unknown period_type: {period_type}" raise ValueError(msg) def period_reset_at( now: datetime, period_type: PeriodType, tz_offset_hours: int = DEFAULT_TZ_OFFSET_HOURS, ) -> datetime | None: """When the current period rolls over (used for the `reset_at` hint).""" if period_type == "lifetime": return None tz = timezone(timedelta(hours=tz_offset_hours)) local = now.astimezone(tz) if period_type == "day": nxt = (local + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0) else: year, month = (local.year + 1, 1) if local.month == 12 else (local.year, local.month + 1) nxt = local.replace( year=year, month=month, day=1, hour=0, minute=0, second=0, microsecond=0 ) return nxt.astimezone(UTC) @dataclass(frozen=True, slots=True) class QuotaState: """Current period counters for one (user, quota_key).""" used: int base_limit: int bonus: int = 0 hard_cap: int | None = None @property def effective_limit(self) -> int: raw = self.base_limit + self.bonus if self.hard_cap is None: return raw return min(raw, self.hard_cap) @property def remaining(self) -> int: return max(self.effective_limit - self.used, 0) @dataclass(frozen=True, slots=True) class QuotaDecision: allowed: bool required: int remaining: int effective_limit: int def evaluate(state: QuotaState, amount: int) -> QuotaDecision: """Server-side hard check (FR-BILL-06). Client hints are UX only.""" if amount <= 0: msg = "amount must be positive" raise ValueError(msg) return QuotaDecision( allowed=state.used + amount <= state.effective_limit, required=amount, remaining=state.remaining, effective_limit=state.effective_limit, ) def reserve(state: QuotaState, amount: int) -> QuotaState: """Freeze ``amount`` units; raises :class:`QuotaExceededError` when over the limit.""" decision = evaluate(state, amount) if not decision.allowed: raise QuotaExceededError(required=amount, remaining=decision.remaining) return replace(state, used=state.used + amount) def refund(state: QuotaState, amount: int) -> QuotaState: """Release a previously reserved amount. Never drives ``used`` below zero.""" if amount <= 0: msg = "amount must be positive" raise ValueError(msg) return replace(state, used=max(state.used - amount, 0)) @dataclass(frozen=True, slots=True) class BonusDecision: """任务奖励发放的封顶判定结果(FR-BILL-03). ``allowed=False`` 时 ``bonus/remaining/effective_limit`` 保持发放前的值: 调用方必须整体放弃,不做部分发放 ——「做任务换 N 次」对用户是明确承诺, 发一半既说不清也无法用流水对账。 """ allowed: bool amount: int bonus: int remaining: int effective_limit: int def grant_bonus(state: QuotaState, amount: int, task_bonus_max: int) -> BonusDecision: """把任务奖励 ``amount`` 计入 ``state.bonus``,超上限则整体拒绝(FR-BILL-03). - 只动 ``bonus``,不动 ``used``:任务奖励抬高的是额度,不是用量; - ``task_bonus_max`` 来自 ``bl_quota_rule``,是任务奖励独立于 ``hard_cap`` 的第二道天花板(两者都生效时 effective_limit 取更紧的); - 边界语义:恰好到顶(bonus + amount == max)允许,超出一格(== max+1)拒绝。 """ if amount <= 0: msg = "amount must be positive" raise ValueError(msg) if state.bonus + amount > task_bonus_max: return BonusDecision( allowed=False, amount=amount, bonus=state.bonus, remaining=state.remaining, effective_limit=state.effective_limit, ) granted = replace(state, bonus=state.bonus + amount) return BonusDecision( allowed=True, amount=amount, bonus=granted.bonus, remaining=granted.remaining, effective_limit=granted.effective_limit, ) def balance_after(state: QuotaState) -> int: """Value written to ``bl_ledger.balance_after``. V1 has no wallet balance, so the ledger records *remaining quota* — that is the number the user actually sees ("今日剩余 N 次"). When V2 introduces ``bl_wallet`` this becomes the wallet balance without a schema change. """ return state.remaining def reserve_idempotency_key(*, user_id: int, quota_key: str, request_key: str) -> str: return f"rsv:{user_id}:{quota_key}:{request_key}" def grant_idempotency_key(*, user_id: int, task_key: str, claim_period_key: str) -> str: """FR-BILL-03 发放幂等:同一任务同一周期派生同一个键. 周期取任务的 ``claim_period``(不是额度行的周期)——「每天可领」的月度 额度任务,幂等粒度必须跟任务承诺走。 """ return f"grant:{user_id}:{task_key}:{claim_period_key}" def refund_idempotency_key(reserve_key: str) -> str: """FR-BILL-05: 退还必须幂等。Derived, so a retry cannot double-refund.""" return f"rfd:{reserve_key}"