fix(billing): unify quota counters, fix Stripe webhooks, tier/plan sync
All checks were successful
Deploy to Production / Build and Deploy (push) Successful in 3m16s
All checks were successful
Deploy to Production / Build and Deploy (push) Successful in 3m16s
This commit is contained in:
@@ -278,6 +278,7 @@ def _db_user_to_model(db_user) -> User:
|
||||
password_hash=db_user.password_hash,
|
||||
avatar_url=db_user.avatar_url,
|
||||
plan=PlanType(db_user.plan) if db_user.plan else PlanType.FREE,
|
||||
tier=db_user.tier or "free",
|
||||
subscription_status=SubscriptionStatus(db_user.subscription_status)
|
||||
if db_user.subscription_status
|
||||
else SubscriptionStatus.ACTIVE,
|
||||
@@ -460,36 +461,48 @@ def update_user(user_id: str, updates: Dict[str, Any]) -> Optional[User]:
|
||||
return User(**users[user_id])
|
||||
|
||||
|
||||
def _reset_usage_if_needed(user_id: str) -> Optional[User]:
|
||||
"""Reset monthly counters if usage_reset_date is in a previous month."""
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
|
||||
with get_sync_session() as session:
|
||||
repo = UserRepository(session)
|
||||
repo.reset_usage_if_needed(user_id)
|
||||
else:
|
||||
user = get_user_by_id(user_id)
|
||||
if not user:
|
||||
return None
|
||||
now = datetime.now(timezone.utc)
|
||||
reset_date = user.usage_reset_date
|
||||
if reset_date.month != now.month or reset_date.year != now.year:
|
||||
update_user(
|
||||
user_id,
|
||||
{
|
||||
"docs_translated_this_month": 0,
|
||||
"pages_translated_this_month": 0,
|
||||
"api_calls_this_month": 0,
|
||||
"usage_reset_date": now.isoformat(),
|
||||
},
|
||||
)
|
||||
return get_user_by_id(user_id)
|
||||
|
||||
|
||||
def check_usage_limits(user: User) -> Dict[str, Any]:
|
||||
"""Check if user has exceeded their plan limits"""
|
||||
# Ensure counters are reset if we've entered a new month.
|
||||
refreshed = _reset_usage_if_needed(user.id)
|
||||
if refreshed:
|
||||
user = refreshed
|
||||
|
||||
plan = PLANS[user.plan]
|
||||
|
||||
# Reset usage if it's a new month
|
||||
now = datetime.now(timezone.utc)
|
||||
if (
|
||||
user.usage_reset_date.month != now.month
|
||||
or user.usage_reset_date.year != now.year
|
||||
):
|
||||
update_user(
|
||||
user.id,
|
||||
{
|
||||
"docs_translated_this_month": 0,
|
||||
"pages_translated_this_month": 0,
|
||||
"api_calls_this_month": 0,
|
||||
"usage_reset_date": now.isoformat() if not USE_DATABASE else now,
|
||||
},
|
||||
)
|
||||
user.docs_translated_this_month = 0
|
||||
user.pages_translated_this_month = 0
|
||||
user.api_calls_this_month = 0
|
||||
|
||||
docs_limit = plan["docs_per_month"]
|
||||
docs_remaining = (
|
||||
max(0, docs_limit - user.docs_translated_this_month) if docs_limit > 0 else -1
|
||||
)
|
||||
unlimited = docs_limit == -1
|
||||
docs_remaining = -1 if unlimited else max(0, docs_limit - user.docs_translated_this_month)
|
||||
|
||||
return {
|
||||
"can_translate": docs_remaining != 0 or user.extra_credits > 0,
|
||||
"can_translate": unlimited or docs_remaining != 0 or user.extra_credits > 0,
|
||||
"docs_used": user.docs_translated_this_month,
|
||||
"docs_limit": docs_limit,
|
||||
"docs_remaining": docs_remaining,
|
||||
@@ -502,23 +515,190 @@ def check_usage_limits(user: User) -> Dict[str, Any]:
|
||||
|
||||
|
||||
def record_usage(
|
||||
user_id: str, pages_count: int, use_credits: bool = False, cost_factor: int = 1
|
||||
user_id: str, pages_count: int, cost_factor: int = 1, reserved_docs: int = 0
|
||||
) -> bool:
|
||||
"""Record document translation usage with optional cost factor depending on AI model"""
|
||||
"""Record document translation usage with optional cost factor depending on AI model.
|
||||
|
||||
`reserved_docs` is the number of document slots already reserved at request time
|
||||
(e.g. by ``reserve_translation_quota``). Those slots are not counted again; only
|
||||
the remaining ``max(0, cost_factor - reserved_docs)`` docs are added here.
|
||||
|
||||
Automatically consumes extra credits first; falls back to monthly quota.
|
||||
Returns True if usage was recorded successfully, False otherwise.
|
||||
"""
|
||||
user = _reset_usage_if_needed(user_id)
|
||||
if not user:
|
||||
return False
|
||||
|
||||
total_cost = pages_count * cost_factor
|
||||
docs_to_record = max(0, cost_factor - reserved_docs)
|
||||
plan = PLANS[user.plan]
|
||||
docs_limit = plan["docs_per_month"]
|
||||
unlimited = docs_limit == -1
|
||||
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
|
||||
with get_sync_session() as session:
|
||||
repo = UserRepository(session)
|
||||
|
||||
if unlimited:
|
||||
# Paid unlimited plans: track usage for analytics/downgrade safety.
|
||||
repo.increment_usage(user_id, docs=docs_to_record, pages=total_cost)
|
||||
return True
|
||||
|
||||
# Prefer credits first, then quota.
|
||||
if user.extra_credits > 0:
|
||||
credits_to_use = min(user.extra_credits, total_cost)
|
||||
if repo.use_credits(user_id, credits_to_use):
|
||||
remaining_cost = total_cost - credits_to_use
|
||||
if remaining_cost > 0:
|
||||
repo.increment_usage(
|
||||
user_id, docs=docs_to_record, pages=remaining_cost
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
if user.docs_translated_this_month + docs_to_record <= docs_limit:
|
||||
repo.increment_usage(user_id, docs=docs_to_record, pages=total_cost)
|
||||
return True
|
||||
return False
|
||||
else:
|
||||
# JSON fallback (non-atomic, dev only)
|
||||
if unlimited:
|
||||
return update_user(
|
||||
user_id,
|
||||
{
|
||||
"docs_translated_this_month": user.docs_translated_this_month
|
||||
+ docs_to_record,
|
||||
"pages_translated_this_month": user.pages_translated_this_month
|
||||
+ total_cost,
|
||||
},
|
||||
) is not None
|
||||
|
||||
if user.extra_credits > 0:
|
||||
credits_to_use = min(user.extra_credits, total_cost)
|
||||
new_credits = user.extra_credits - credits_to_use
|
||||
remaining_cost = total_cost - credits_to_use
|
||||
updates = {"extra_credits": new_credits}
|
||||
if remaining_cost > 0:
|
||||
if user.docs_translated_this_month + docs_to_record > docs_limit:
|
||||
return False
|
||||
updates["docs_translated_this_month"] = (
|
||||
user.docs_translated_this_month + docs_to_record
|
||||
)
|
||||
updates["pages_translated_this_month"] = (
|
||||
user.pages_translated_this_month + remaining_cost
|
||||
)
|
||||
return update_user(user_id, updates) is not None
|
||||
|
||||
if user.docs_translated_this_month + docs_to_record <= docs_limit:
|
||||
return update_user(
|
||||
user_id,
|
||||
{
|
||||
"docs_translated_this_month": user.docs_translated_this_month
|
||||
+ docs_to_record,
|
||||
"pages_translated_this_month": user.pages_translated_this_month
|
||||
+ total_cost,
|
||||
},
|
||||
) is not None
|
||||
return False
|
||||
|
||||
|
||||
def reserve_translation_quota(user_id: str) -> bool:
|
||||
"""Atomically reserve one document slot at request time.
|
||||
|
||||
Returns True if the reservation succeeded. This prevents race conditions where
|
||||
multiple concurrent requests could exceed the monthly quota.
|
||||
"""
|
||||
user = _reset_usage_if_needed(user_id)
|
||||
if not user:
|
||||
return False
|
||||
|
||||
plan = PLANS[user.plan]
|
||||
docs_limit = plan["docs_per_month"]
|
||||
unlimited = docs_limit == -1
|
||||
|
||||
if unlimited:
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
|
||||
with get_sync_session() as session:
|
||||
repo = UserRepository(session)
|
||||
repo.increment_usage(user_id, docs=1, pages=0)
|
||||
return True
|
||||
return update_user(
|
||||
user_id,
|
||||
{"docs_translated_this_month": user.docs_translated_this_month + 1},
|
||||
) is not None
|
||||
|
||||
# Limited plan: prefer credits first, then quota.
|
||||
if user.extra_credits > 0:
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
|
||||
with get_sync_session() as session:
|
||||
repo = UserRepository(session)
|
||||
return repo.use_credits(user_id, 1)
|
||||
return update_user(
|
||||
user_id, {"extra_credits": max(0, user.extra_credits - 1)}
|
||||
) is not None
|
||||
|
||||
if user.docs_translated_this_month + 1 <= docs_limit:
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
|
||||
with get_sync_session() as session:
|
||||
repo = UserRepository(session)
|
||||
repo.increment_usage(user_id, docs=1, pages=0)
|
||||
return True
|
||||
return update_user(
|
||||
user_id,
|
||||
{"docs_translated_this_month": user.docs_translated_this_month + 1},
|
||||
) is not None
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def release_translation_quota(user_id: str) -> bool:
|
||||
"""Release a previously reserved document slot (e.g. on translation failure)."""
|
||||
user = get_user_by_id(user_id)
|
||||
if not user:
|
||||
return False
|
||||
|
||||
updates = {
|
||||
"docs_translated_this_month": user.docs_translated_this_month + cost_factor,
|
||||
"pages_translated_this_month": user.pages_translated_this_month + (pages_count * cost_factor),
|
||||
}
|
||||
if USE_DATABASE and DATABASE_AVAILABLE:
|
||||
from database.connection import get_sync_session
|
||||
from database.repositories import UserRepository
|
||||
from sqlalchemy import update
|
||||
|
||||
if use_credits:
|
||||
updates["extra_credits"] = max(0, user.extra_credits - (pages_count * cost_factor))
|
||||
with get_sync_session() as session:
|
||||
now = datetime.now(timezone.utc)
|
||||
session.execute(
|
||||
update(db_models.User)
|
||||
.where(
|
||||
db_models.User.id == user_id,
|
||||
db_models.User.docs_translated_this_month > 0,
|
||||
)
|
||||
.values(
|
||||
docs_translated_this_month=db_models.User.docs_translated_this_month
|
||||
- 1,
|
||||
updated_at=now,
|
||||
)
|
||||
.execution_options(synchronize_session=False)
|
||||
)
|
||||
session.commit()
|
||||
return True
|
||||
|
||||
result = update_user(user_id, updates)
|
||||
return result is not None
|
||||
if user.docs_translated_this_month > 0:
|
||||
return update_user(
|
||||
user_id,
|
||||
{"docs_translated_this_month": user.docs_translated_this_month - 1},
|
||||
) is not None
|
||||
return True
|
||||
|
||||
|
||||
def add_credits(user_id: str, credits: int) -> bool:
|
||||
|
||||
Reference in New Issue
Block a user