Spaces:

asnannp
/

DocDoeAI

Runtime error

DocDoeAI / app /services /operation_usage.py

deploy: sync backend to Space root (learn-lesson HF cache fix)

65d68f5 about 1 month ago

19.3 kB

"""Authoritative request-unit limits for costly DocDoe operations.
Policy is intentionally simple: a successful operation consumes one or more
student-readable units; rejection and provider failure do not. Reservations
count toward the limit and concurrency cap while work is running. A logical
retry reuses the same user/operation/idempotency row, so it cannot double-count.
Prompts and response bodies are never stored here.
"""
from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Literal

from fastapi import HTTPException, status
from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core.config import get_settings
from app.models.usage_operation import UsageOperation
from app.models.user import User
from app.services.entitlements import EntitlementFeature, require_entitlement_feature, resolve_entitlements
from app.services.plan_catalog import get_monthly_usage_limits, get_plan_config
from app.utils.request_id import request_id_ctx_var

logger = logging.getLogger(__name__)

UsageOperationName = Literal[
    "study_chat_ai",
    "tuition_tutor",
    "correction_ai",
    "source_processing",
    "learn_anything_generation",
    "pyq_style_generation",
    "video_generation",
    "upload_storage",
    "exam_intelligence_generation",
]

_DEFAULT_POLICY: dict[str, dict[str, int | str]] = {
    "study_chat_ai": {"limit": 50, "period": "day", "max_active": 1},
    "tuition_tutor": {"limit": 20, "period": "day", "max_active": 1},
    "correction_ai": {"limit": 20, "period": "day", "max_active": 1},
    "source_processing": {"limit": 10, "period": "month", "max_active": 1},
    "learn_anything_generation": {"limit": 10, "period": "day", "max_active": 1},
    "pyq_style_generation": {"limit": 20, "period": "day", "max_active": 1},
    "video_generation": {"limit": 2, "period": "month", "max_active": 1},
    "upload_storage": {"limit": 100 * 1024 * 1024, "period": "month", "max_active": 2},
    # Audit H3: the two generative GET endpoints in /intelligence/ previously
    # had no per-user quota at all. They share the exam-intelligence feature
    # gate with pyq_style_generation but get their own daily budget.
    "exam_intelligence_generation": {"limit": 20, "period": "day", "max_active": 1},
}

_OPERATION_FEATURES: dict[str, EntitlementFeature] = {
    "study_chat_ai": "study_chat_ai",
    "tuition_tutor": "tuition_ai_tutor",
    "correction_ai": "deterministic_correction",
    "source_processing": "source_upload",
    "learn_anything_generation": "learn_anything_generation",
    "pyq_style_generation": "pyq_style_generation",
    "video_generation": "heavy_generation",
    "upload_storage": "source_upload",
    "exam_intelligence_generation": "pyq_style_generation",
}

# Estimated provider USD cost per AI-token operation, used ONLY for the
# per-user daily budget in reserve_usage(). Values are derived from the
# provider cost table and typical token volumes, not billed amounts. Operations
# that do not consume LLM tokens (source_processing, video_generation,
# upload_storage) are intentionally absent so they are never double-counted.
_OPERATION_ESTIMATED_COST_USD: dict[str, float] = {
    "study_chat_ai": 0.002,
    "tuition_tutor": 0.004,
    "correction_ai": 0.003,
    "learn_anything_generation": 0.008,
    "pyq_style_generation": 0.006,
    "exam_intelligence_generation": 0.006,
}


@dataclass(frozen=True)
class UsagePolicy:
    limit: int
    period: Literal["day", "month"]
    max_active: int


@dataclass(frozen=True)
class UsageReservation:
    record: UsageOperation
    duplicate_completed: bool = False

    @property
    def should_charge(self) -> bool:
        return not self.duplicate_completed


def _policies(policy_key: str = "free_trial") -> dict[str, UsagePolicy]:
    raw: dict[str, Any] = {}
    configured = get_settings().usage_policy_json.strip()
    if configured:
        try:
            parsed = json.loads(configured)
            if isinstance(parsed, dict):
                raw = parsed
        except (TypeError, ValueError):
            logger.warning("usage_policy_invalid", extra={"event": "usage_policy_invalid"})
    free_config = get_plan_config("free_trial")
    plan_config = get_plan_config(policy_key) or free_config
    generation_multiplier = (
        max(1, plan_config.monthly_generation_limit // max(1, free_config.monthly_generation_limit))
        if plan_config and free_config
        else 1
    )
    legacy_limits = get_monthly_usage_limits(policy_key)
    policy_defaults = {name: dict(values) for name, values in _DEFAULT_POLICY.items()}
    for operation in (
        "study_chat_ai",
        "tuition_tutor",
        "correction_ai",
        "learn_anything_generation",
        "pyq_style_generation",
        "exam_intelligence_generation",
    ):
        policy_defaults[operation]["limit"] = int(policy_defaults[operation]["limit"]) * generation_multiplier
    policy_defaults["source_processing"]["limit"] = int(
        legacy_limits.get("documents_uploaded_count", policy_defaults["source_processing"]["limit"])
    )
    policy_defaults["video_generation"]["limit"] = int(
        plan_config.monthly_video_limit if plan_config else policy_defaults["video_generation"]["limit"]
    )
    policy_defaults["upload_storage"]["limit"] = int(
        legacy_limits.get("storage_bytes_used", policy_defaults["upload_storage"]["limit"])
    )
    policy_overrides = raw.get("policies", {}) if isinstance(raw.get("policies"), dict) else {}
    selected_overrides = (
        policy_overrides.get(policy_key, {})
        if isinstance(policy_overrides.get(policy_key), dict)
        else {}
    )
    result: dict[str, UsagePolicy] = {}
    for operation, defaults in policy_defaults.items():
        global_override = raw.get(operation) if isinstance(raw.get(operation), dict) else {}
        specific_override = (
            selected_overrides.get(operation)
            if isinstance(selected_overrides.get(operation), dict)
            else {}
        )
        override = {**global_override, **specific_override}
        period = override.get("period", defaults["period"])
        result[operation] = UsagePolicy(
            limit=max(0, int(override.get("limit", defaults["limit"]))),
            period="month" if period == "month" else "day",
            max_active=max(1, int(override.get("max_active", defaults["max_active"]))),
        )
    return result


def _period_key(policy: UsagePolicy, now: datetime) -> str:
    return now.strftime("%Y-%m" if policy.period == "month" else "%Y-%m-%d")


def stable_idempotency_key(*parts: object) -> str:
    """Hash metadata into a bounded key without persisting private prompt text."""
    raw = "\x1f".join(str(part) for part in parts)
    return hashlib.sha256(raw.encode("utf-8", errors="ignore")).hexdigest()


def _student_error(code: str, message: str, *, http_status: int) -> HTTPException:
    return HTTPException(
        status_code=http_status,
        detail={
            "code": code,
            "message": message,
            "retryable": code == "OPERATION_IN_PROGRESS",
            "request_id": request_id_ctx_var.get("-"),
        },
    )


def _duplicate_completed_error() -> HTTPException:
    """Reject a replay of an already-completed idempotency key.
    Used by callers whose completed row stores no result: replaying the key
    would otherwise re-run the paid operation for free.
    """
    return _student_error(
        "DUPLICATE_REQUEST",
        "This request has already been completed. Start a new request to continue.",
        http_status=status.HTTP_409_CONFLICT,
    )


def _log(event: str, *, user_id: str, operation: str, units: int, policy: UsagePolicy, remaining: int) -> None:
    logger.info(
        event,
        extra={
            "event": event,
            "operation": operation,
            "user_ref": stable_idempotency_key(user_id)[:16],
            "units": units,
            "remaining": max(0, remaining),
            "limit": policy.limit,
            "period": policy.period,
        },
    )


def _daily_budget_denial(
    db: Session, *, user_id: str, operation: str, now: datetime
) -> tuple[float, float, float] | None:
    """Return (day_cost, this_cost, cap) when the per-user daily AI budget is
    exceeded, else None.
    Uses the same day period_key convention and reserved/completed statuses as
    the unit-limit query above. Fails OPEN: any accounting error returns None so
    a broken cost lookup can never take down AI generation for students.
    """
    try:
        cap = float(get_settings().ai_router_per_user_daily_budget_usd or 0)
        if cap <= 0:
            return None
        cost = float(_OPERATION_ESTIMATED_COST_USD.get(operation, 0) or 0)
        if cost <= 0:
            return None
        day_key = now.strftime("%Y-%m-%d")
        rows = db.execute(
            select(UsageOperation.operation, func.coalesce(func.sum(UsageOperation.units), 0))
            .where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation.in_(list(_OPERATION_ESTIMATED_COST_USD)),
                UsageOperation.period_key == day_key,
                UsageOperation.status.in_(["reserved", "completed"]),
            )
            .group_by(UsageOperation.operation)
        ).all()
        day_cost = sum(
            float(_OPERATION_ESTIMATED_COST_USD.get(op, 0)) * int(units or 0)
            for op, units in rows
        )
        if day_cost + cost > cap:
            return day_cost, cost, cap
        return None
    except Exception:  # noqa: BLE001 - accounting must fail open
        logger.warning(
            "usage_budget_check_failed",
            exc_info=True,
            extra={"event": "usage_budget_check_failed", "operation": operation},
        )
        return None


def reserve_usage(
    db: Session,
    *,
    user_id: str,
    operation: UsageOperationName,
    idempotency_key: str,
    units: int = 1,
    reject_completed_duplicate: bool = False,
) -> UsageReservation:
    """Reserve quota before work.
    A completed row whose idempotency key is reused is normally returned as a
    no-charge duplicate. That is correct for resource-backed operations
    (uploads, video renders) which look up the stored ``resource_id`` and return
    the existing artefact. Callers whose completed row stores no result — chat,
    where replaying the key would otherwise re-run a paid AI call for free —
    pass ``reject_completed_duplicate=True`` to reject the replay instead.
    """
    if units <= 0:
        raise ValueError("Usage units must be positive")
    entitlement = require_entitlement_feature(db, user_id, _OPERATION_FEATURES[operation])
    policy = _policies(entitlement.usage_policy_key)[operation]
    if not get_settings().usage_limits_enabled:
        policy = UsagePolicy(limit=2_147_483_647, period=policy.period, max_active=policy.max_active)
    now = datetime.now(timezone.utc)
    period_key = _period_key(policy, now)
    key = idempotency_key.strip()[:180] or stable_idempotency_key(request_id_ctx_var.get("-"))

    existing = db.scalar(
        select(UsageOperation).where(
            UsageOperation.user_id == user_id,
            UsageOperation.operation == operation,
            UsageOperation.idempotency_key == key,
        )
    )
    if existing is not None and existing.status == "completed":
        if reject_completed_duplicate:
            _log("usage_duplicate_rejected", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
            raise _duplicate_completed_error()
        _log("usage_allowed", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
        return UsageReservation(existing, duplicate_completed=True)
    if existing is not None and existing.status == "reserved":
        cutoff = now - timedelta(seconds=get_settings().usage_reservation_ttl_seconds)
        updated_at = existing.updated_at
        if updated_at.tzinfo is None:
            updated_at = updated_at.replace(tzinfo=timezone.utc)
        if updated_at >= cutoff:
            raise _student_error(
                "OPERATION_IN_PROGRESS",
                "That request is already running.",
                http_status=status.HTTP_409_CONFLICT,
            )
        existing.status = "released"
        db.add(existing)
        db.commit()

    # PostgreSQL serializes reservations per user. SQLite remains deterministic
    # in local tests through its write lock plus the unique logical key.
    db.scalar(select(User.id).where(User.id == user_id).with_for_update())
    used = int(
        db.scalar(
            select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.period_key == period_key,
                UsageOperation.status.in_(["reserved", "completed"]),
            )
        )
        or 0
    )
    if used + units > policy.limit:
        _log("usage_denied", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used)
        raise _student_error(
            "USAGE_LIMIT_REACHED",
            "You've reached today's AI study limit. Try again when your allowance resets.",
            http_status=status.HTTP_402_PAYMENT_REQUIRED,
        )

    # Per-user daily USD budget: a cost backstop on top of the count limits
    # above. Denies through the same path/shape as the unit limit.
    budget_denial = _daily_budget_denial(db, user_id=user_id, operation=operation, now=now)
    if budget_denial is not None:
        day_cost, operation_cost, cap = budget_denial
        logger.warning(
            "usage_budget_denied",
            extra={
                "event": "usage_budget_denied",
                "operation": operation,
                "user_ref": stable_idempotency_key(user_id)[:16],
                "estimated_day_cost_usd": round(day_cost, 6),
                "estimated_operation_cost_usd": round(operation_cost, 6),
                "cap_usd": cap,
            },
        )
        raise _student_error(
            "USAGE_LIMIT_REACHED",
            "You've reached today's AI study budget. Try again when your allowance resets.",
            http_status=status.HTTP_402_PAYMENT_REQUIRED,
        )

    active = int(
        db.scalar(
            select(func.count()).select_from(UsageOperation).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.status == "reserved",
            )
        )
        or 0
    )
    if active >= policy.max_active:
        raise _student_error(
            "OPERATION_IN_PROGRESS",
            "That request is already running.",
            http_status=status.HTTP_409_CONFLICT,
        )

    if existing is not None:
        existing.status = "reserved"
        existing.period_key = period_key
        existing.units = units
        existing.request_id = request_id_ctx_var.get("-")
        existing.completed_at = None
        record = existing
    else:
        record = UsageOperation(
            user_id=user_id,
            operation=operation,
            period_key=period_key,
            idempotency_key=key,
            units=units,
            status="reserved",
            request_id=request_id_ctx_var.get("-"),
        )
    db.add(record)
    try:
        db.commit()
        db.refresh(record)
    except IntegrityError:
        db.rollback()
        raced = db.scalar(
            select(UsageOperation).where(
                UsageOperation.user_id == user_id,
                UsageOperation.operation == operation,
                UsageOperation.idempotency_key == key,
            )
        )
        if raced is not None and raced.status == "completed":
            if reject_completed_duplicate:
                raise _duplicate_completed_error()
            return UsageReservation(raced, duplicate_completed=True)
        raise _student_error(
            "OPERATION_IN_PROGRESS",
            "That request is already running.",
            http_status=status.HTTP_409_CONFLICT,
        )
    _log("usage_allowed", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used - units)
    return UsageReservation(record)


def finalize_usage(db: Session, reservation: UsageReservation, *, resource_id: str | None = None) -> None:
    if not reservation.should_charge:
        return
    db.execute(
        update(UsageOperation)
        .where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
        .values(status="completed", completed_at=datetime.now(timezone.utc), resource_id=resource_id)
    )
    db.commit()
    entitlement = resolve_entitlements(db, reservation.record.user_id)
    policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
    _log("usage_recorded", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


def release_usage(db: Session, reservation: UsageReservation) -> None:
    if not reservation.should_charge:
        return
    db.execute(
        update(UsageOperation)
        .where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
        .values(status="released", completed_at=None)
    )
    db.commit()
    entitlement = resolve_entitlements(db, reservation.record.user_id)
    policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
    _log("usage_reservation_released", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


def get_usage_status(db: Session, user_id: str) -> dict[str, Any]:
    entitlement = resolve_entitlements(db, user_id)
    items: list[dict[str, Any]] = []
    for operation, policy in _policies(entitlement.usage_policy_key).items():
        period_key = _period_key(policy, datetime.now(timezone.utc))
        used = int(
            db.scalar(
                select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
                    UsageOperation.user_id == user_id,
                    UsageOperation.operation == operation,
                    UsageOperation.period_key == period_key,
                    UsageOperation.status.in_(["reserved", "completed"]),
                )
            )
            or 0
        )
        items.append(
            {
                "operation": operation,
                "period": policy.period,
                "used": used,
                "limit": policy.limit,
                "remaining": max(0, policy.limit - used),
                "allowed": used < policy.limit,
            }
        )
    return {
        "tier": entitlement.tier,
        "entitlement_status": entitlement.status,
        "usage_policy_key": entitlement.usage_policy_key,
        "items": items,
    }

	"""Authoritative request-unit limits for costly DocDoe operations.

	Policy is intentionally simple: a successful operation consumes one or more
	student-readable units; rejection and provider failure do not. Reservations
	count toward the limit and concurrency cap while work is running. A logical
	retry reuses the same user/operation/idempotency row, so it cannot double-count.
	Prompts and response bodies are never stored here.
	"""
	from __future__ import annotations

	import hashlib
	import json
	import logging
	from dataclasses import dataclass
	from datetime import datetime, timedelta, timezone
	from typing import Any, Literal

	from fastapi import HTTPException, status
	from sqlalchemy import func, select, update
	from sqlalchemy.exc import IntegrityError
	from sqlalchemy.orm import Session

	from app.core.config import get_settings
	from app.models.usage_operation import UsageOperation
	from app.models.user import User
	from app.services.entitlements import EntitlementFeature, require_entitlement_feature, resolve_entitlements
	from app.services.plan_catalog import get_monthly_usage_limits, get_plan_config
	from app.utils.request_id import request_id_ctx_var

	logger = logging.getLogger(__name__)

	UsageOperationName = Literal[
	"study_chat_ai",
	"tuition_tutor",
	"correction_ai",
	"source_processing",
	"learn_anything_generation",
	"pyq_style_generation",
	"video_generation",
	"upload_storage",
	"exam_intelligence_generation",
	]

	_DEFAULT_POLICY: dict[str, dict[str, int \| str]] = {
	"study_chat_ai": {"limit": 50, "period": "day", "max_active": 1},
	"tuition_tutor": {"limit": 20, "period": "day", "max_active": 1},
	"correction_ai": {"limit": 20, "period": "day", "max_active": 1},
	"source_processing": {"limit": 10, "period": "month", "max_active": 1},
	"learn_anything_generation": {"limit": 10, "period": "day", "max_active": 1},
	"pyq_style_generation": {"limit": 20, "period": "day", "max_active": 1},
	"video_generation": {"limit": 2, "period": "month", "max_active": 1},
	"upload_storage": {"limit": 100 * 1024 * 1024, "period": "month", "max_active": 2},
	# Audit H3: the two generative GET endpoints in /intelligence/ previously
	# had no per-user quota at all. They share the exam-intelligence feature
	# gate with pyq_style_generation but get their own daily budget.
	"exam_intelligence_generation": {"limit": 20, "period": "day", "max_active": 1},
	}

	_OPERATION_FEATURES: dict[str, EntitlementFeature] = {
	"study_chat_ai": "study_chat_ai",
	"tuition_tutor": "tuition_ai_tutor",
	"correction_ai": "deterministic_correction",
	"source_processing": "source_upload",
	"learn_anything_generation": "learn_anything_generation",
	"pyq_style_generation": "pyq_style_generation",
	"video_generation": "heavy_generation",
	"upload_storage": "source_upload",
	"exam_intelligence_generation": "pyq_style_generation",
	}

	# Estimated provider USD cost per AI-token operation, used ONLY for the
	# per-user daily budget in reserve_usage(). Values are derived from the
	# provider cost table and typical token volumes, not billed amounts. Operations
	# that do not consume LLM tokens (source_processing, video_generation,
	# upload_storage) are intentionally absent so they are never double-counted.
	_OPERATION_ESTIMATED_COST_USD: dict[str, float] = {
	"study_chat_ai": 0.002,
	"tuition_tutor": 0.004,
	"correction_ai": 0.003,
	"learn_anything_generation": 0.008,
	"pyq_style_generation": 0.006,
	"exam_intelligence_generation": 0.006,
	}


	@dataclass(frozen=True)
	class UsagePolicy:
	limit: int
	period: Literal["day", "month"]
	max_active: int


	@dataclass(frozen=True)
	class UsageReservation:
	record: UsageOperation
	duplicate_completed: bool = False

	@property
	def should_charge(self) -> bool:
	return not self.duplicate_completed


	def _policies(policy_key: str = "free_trial") -> dict[str, UsagePolicy]:
	raw: dict[str, Any] = {}
	configured = get_settings().usage_policy_json.strip()
	if configured:
	try:
	parsed = json.loads(configured)
	if isinstance(parsed, dict):
	raw = parsed
	except (TypeError, ValueError):
	logger.warning("usage_policy_invalid", extra={"event": "usage_policy_invalid"})
	free_config = get_plan_config("free_trial")
	plan_config = get_plan_config(policy_key) or free_config
	generation_multiplier = (
	max(1, plan_config.monthly_generation_limit // max(1, free_config.monthly_generation_limit))
	if plan_config and free_config
	else 1
	)
	legacy_limits = get_monthly_usage_limits(policy_key)
	policy_defaults = {name: dict(values) for name, values in _DEFAULT_POLICY.items()}
	for operation in (
	"study_chat_ai",
	"tuition_tutor",
	"correction_ai",
	"learn_anything_generation",
	"pyq_style_generation",
	"exam_intelligence_generation",
	):
	policy_defaults[operation]["limit"] = int(policy_defaults[operation]["limit"]) * generation_multiplier
	policy_defaults["source_processing"]["limit"] = int(
	legacy_limits.get("documents_uploaded_count", policy_defaults["source_processing"]["limit"])
	)
	policy_defaults["video_generation"]["limit"] = int(
	plan_config.monthly_video_limit if plan_config else policy_defaults["video_generation"]["limit"]
	)
	policy_defaults["upload_storage"]["limit"] = int(
	legacy_limits.get("storage_bytes_used", policy_defaults["upload_storage"]["limit"])
	)
	policy_overrides = raw.get("policies", {}) if isinstance(raw.get("policies"), dict) else {}
	selected_overrides = (
	policy_overrides.get(policy_key, {})
	if isinstance(policy_overrides.get(policy_key), dict)
	else {}
	)
	result: dict[str, UsagePolicy] = {}
	for operation, defaults in policy_defaults.items():
	global_override = raw.get(operation) if isinstance(raw.get(operation), dict) else {}
	specific_override = (
	selected_overrides.get(operation)
	if isinstance(selected_overrides.get(operation), dict)
	else {}
	)
	override = {global_override, specific_override}
	period = override.get("period", defaults["period"])
	result[operation] = UsagePolicy(
	limit=max(0, int(override.get("limit", defaults["limit"]))),
	period="month" if period == "month" else "day",
	max_active=max(1, int(override.get("max_active", defaults["max_active"]))),
	)
	return result


	def _period_key(policy: UsagePolicy, now: datetime) -> str:
	return now.strftime("%Y-%m" if policy.period == "month" else "%Y-%m-%d")


	def stable_idempotency_key(*parts: object) -> str:
	"""Hash metadata into a bounded key without persisting private prompt text."""
	raw = "\x1f".join(str(part) for part in parts)
	return hashlib.sha256(raw.encode("utf-8", errors="ignore")).hexdigest()


	def _student_error(code: str, message: str, *, http_status: int) -> HTTPException:
	return HTTPException(
	status_code=http_status,
	detail={
	"code": code,
	"message": message,
	"retryable": code == "OPERATION_IN_PROGRESS",
	"request_id": request_id_ctx_var.get("-"),
	},
	)


	def _duplicate_completed_error() -> HTTPException:
	"""Reject a replay of an already-completed idempotency key.

	Used by callers whose completed row stores no result: replaying the key
	would otherwise re-run the paid operation for free.
	"""
	return _student_error(
	"DUPLICATE_REQUEST",
	"This request has already been completed. Start a new request to continue.",
	http_status=status.HTTP_409_CONFLICT,
	)


	def _log(event: str, *, user_id: str, operation: str, units: int, policy: UsagePolicy, remaining: int) -> None:
	logger.info(
	event,
	extra={
	"event": event,
	"operation": operation,
	"user_ref": stable_idempotency_key(user_id)[:16],
	"units": units,
	"remaining": max(0, remaining),
	"limit": policy.limit,
	"period": policy.period,
	},
	)


	def _daily_budget_denial(
	db: Session, *, user_id: str, operation: str, now: datetime
	) -> tuple[float, float, float] \| None:
	"""Return (day_cost, this_cost, cap) when the per-user daily AI budget is
	exceeded, else None.

	Uses the same day period_key convention and reserved/completed statuses as
	the unit-limit query above. Fails OPEN: any accounting error returns None so
	a broken cost lookup can never take down AI generation for students.
	"""
	try:
	cap = float(get_settings().ai_router_per_user_daily_budget_usd or 0)
	if cap <= 0:
	return None
	cost = float(_OPERATION_ESTIMATED_COST_USD.get(operation, 0) or 0)
	if cost <= 0:
	return None
	day_key = now.strftime("%Y-%m-%d")
	rows = db.execute(
	select(UsageOperation.operation, func.coalesce(func.sum(UsageOperation.units), 0))
	.where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation.in_(list(_OPERATION_ESTIMATED_COST_USD)),
	UsageOperation.period_key == day_key,
	UsageOperation.status.in_(["reserved", "completed"]),
	)
	.group_by(UsageOperation.operation)
	).all()
	day_cost = sum(
	float(_OPERATION_ESTIMATED_COST_USD.get(op, 0)) * int(units or 0)
	for op, units in rows
	)
	if day_cost + cost > cap:
	return day_cost, cost, cap
	return None
	except Exception: # noqa: BLE001 - accounting must fail open
	logger.warning(
	"usage_budget_check_failed",
	exc_info=True,
	extra={"event": "usage_budget_check_failed", "operation": operation},
	)
	return None


	def reserve_usage(
	db: Session,
	*,
	user_id: str,
	operation: UsageOperationName,
	idempotency_key: str,
	units: int = 1,
	reject_completed_duplicate: bool = False,
	) -> UsageReservation:
	"""Reserve quota before work.

	A completed row whose idempotency key is reused is normally returned as a
	no-charge duplicate. That is correct for resource-backed operations
	(uploads, video renders) which look up the stored ``resource_id`` and return
	the existing artefact. Callers whose completed row stores no result — chat,
	where replaying the key would otherwise re-run a paid AI call for free —
	pass ``reject_completed_duplicate=True`` to reject the replay instead.
	"""
	if units <= 0:
	raise ValueError("Usage units must be positive")
	entitlement = require_entitlement_feature(db, user_id, _OPERATION_FEATURES[operation])
	policy = _policies(entitlement.usage_policy_key)[operation]
	if not get_settings().usage_limits_enabled:
	policy = UsagePolicy(limit=2_147_483_647, period=policy.period, max_active=policy.max_active)
	now = datetime.now(timezone.utc)
	period_key = _period_key(policy, now)
	key = idempotency_key.strip()[:180] or stable_idempotency_key(request_id_ctx_var.get("-"))

	existing = db.scalar(
	select(UsageOperation).where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation == operation,
	UsageOperation.idempotency_key == key,
	)
	)
	if existing is not None and existing.status == "completed":
	if reject_completed_duplicate:
	_log("usage_duplicate_rejected", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
	raise _duplicate_completed_error()
	_log("usage_allowed", user_id=user_id, operation=operation, units=0, policy=policy, remaining=policy.limit)
	return UsageReservation(existing, duplicate_completed=True)
	if existing is not None and existing.status == "reserved":
	cutoff = now - timedelta(seconds=get_settings().usage_reservation_ttl_seconds)
	updated_at = existing.updated_at
	if updated_at.tzinfo is None:
	updated_at = updated_at.replace(tzinfo=timezone.utc)
	if updated_at >= cutoff:
	raise _student_error(
	"OPERATION_IN_PROGRESS",
	"That request is already running.",
	http_status=status.HTTP_409_CONFLICT,
	)
	existing.status = "released"
	db.add(existing)
	db.commit()

	# PostgreSQL serializes reservations per user. SQLite remains deterministic
	# in local tests through its write lock plus the unique logical key.
	db.scalar(select(User.id).where(User.id == user_id).with_for_update())
	used = int(
	db.scalar(
	select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation == operation,
	UsageOperation.period_key == period_key,
	UsageOperation.status.in_(["reserved", "completed"]),
	)
	)
	or 0
	)
	if used + units > policy.limit:
	_log("usage_denied", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used)
	raise _student_error(
	"USAGE_LIMIT_REACHED",
	"You've reached today's AI study limit. Try again when your allowance resets.",
	http_status=status.HTTP_402_PAYMENT_REQUIRED,
	)

	# Per-user daily USD budget: a cost backstop on top of the count limits
	# above. Denies through the same path/shape as the unit limit.
	budget_denial = _daily_budget_denial(db, user_id=user_id, operation=operation, now=now)
	if budget_denial is not None:
	day_cost, operation_cost, cap = budget_denial
	logger.warning(
	"usage_budget_denied",
	extra={
	"event": "usage_budget_denied",
	"operation": operation,
	"user_ref": stable_idempotency_key(user_id)[:16],
	"estimated_day_cost_usd": round(day_cost, 6),
	"estimated_operation_cost_usd": round(operation_cost, 6),
	"cap_usd": cap,
	},
	)
	raise _student_error(
	"USAGE_LIMIT_REACHED",
	"You've reached today's AI study budget. Try again when your allowance resets.",
	http_status=status.HTTP_402_PAYMENT_REQUIRED,
	)

	active = int(
	db.scalar(
	select(func.count()).select_from(UsageOperation).where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation == operation,
	UsageOperation.status == "reserved",
	)
	)
	or 0
	)
	if active >= policy.max_active:
	raise _student_error(
	"OPERATION_IN_PROGRESS",
	"That request is already running.",
	http_status=status.HTTP_409_CONFLICT,
	)

	if existing is not None:
	existing.status = "reserved"
	existing.period_key = period_key
	existing.units = units
	existing.request_id = request_id_ctx_var.get("-")
	existing.completed_at = None
	record = existing
	else:
	record = UsageOperation(
	user_id=user_id,
	operation=operation,
	period_key=period_key,
	idempotency_key=key,
	units=units,
	status="reserved",
	request_id=request_id_ctx_var.get("-"),
	)
	db.add(record)
	try:
	db.commit()
	db.refresh(record)
	except IntegrityError:
	db.rollback()
	raced = db.scalar(
	select(UsageOperation).where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation == operation,
	UsageOperation.idempotency_key == key,
	)
	)
	if raced is not None and raced.status == "completed":
	if reject_completed_duplicate:
	raise _duplicate_completed_error()
	return UsageReservation(raced, duplicate_completed=True)
	raise _student_error(
	"OPERATION_IN_PROGRESS",
	"That request is already running.",
	http_status=status.HTTP_409_CONFLICT,
	)
	_log("usage_allowed", user_id=user_id, operation=operation, units=units, policy=policy, remaining=policy.limit - used - units)
	return UsageReservation(record)


	def finalize_usage(db: Session, reservation: UsageReservation, *, resource_id: str \| None = None) -> None:
	if not reservation.should_charge:
	return
	db.execute(
	update(UsageOperation)
	.where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
	.values(status="completed", completed_at=datetime.now(timezone.utc), resource_id=resource_id)
	)
	db.commit()
	entitlement = resolve_entitlements(db, reservation.record.user_id)
	policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
	_log("usage_recorded", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


	def release_usage(db: Session, reservation: UsageReservation) -> None:
	if not reservation.should_charge:
	return
	db.execute(
	update(UsageOperation)
	.where(UsageOperation.id == reservation.record.id, UsageOperation.status == "reserved")
	.values(status="released", completed_at=None)
	)
	db.commit()
	entitlement = resolve_entitlements(db, reservation.record.user_id)
	policy = _policies(entitlement.usage_policy_key)[reservation.record.operation]
	_log("usage_reservation_released", user_id=reservation.record.user_id, operation=reservation.record.operation, units=reservation.record.units, policy=policy, remaining=policy.limit)


	def get_usage_status(db: Session, user_id: str) -> dict[str, Any]:
	entitlement = resolve_entitlements(db, user_id)
	items: list[dict[str, Any]] = []
	for operation, policy in _policies(entitlement.usage_policy_key).items():
	period_key = _period_key(policy, datetime.now(timezone.utc))
	used = int(
	db.scalar(
	select(func.coalesce(func.sum(UsageOperation.units), 0)).where(
	UsageOperation.user_id == user_id,
	UsageOperation.operation == operation,
	UsageOperation.period_key == period_key,
	UsageOperation.status.in_(["reserved", "completed"]),
	)
	)
	or 0
	)
	items.append(
	{
	"operation": operation,
	"period": policy.period,
	"used": used,
	"limit": policy.limit,
	"remaining": max(0, policy.limit - used),
	"allowed": used < policy.limit,
	}
	)
	return {
	"tier": entitlement.tier,
	"entitlement_status": entitlement.status,
	"usage_policy_key": entitlement.usage_policy_key,
	"items": items,
	}