From dc76ca10c7b5ec9a5c23618f88c71d970e49a107 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 9 May 2026 23:31:46 +0800 Subject: [PATCH] fix: initialize missing user balance during tier lookup --- .../app/services/rate_limit/tier_service.py | 24 ++++- .../tests/contract/test_billing_contract.py | 94 +++++++++++++++++++ .../services/billing/credits_service.py | 9 +- 3 files changed, 121 insertions(+), 6 deletions(-) diff --git a/apps/api/app/services/rate_limit/tier_service.py b/apps/api/app/services/rate_limit/tier_service.py index 01664e922..91bb73d00 100644 --- a/apps/api/app/services/rate_limit/tier_service.py +++ b/apps/api/app/services/rate_limit/tier_service.py @@ -20,6 +20,7 @@ from shared.models.database.tier_limit import TierLimit from shared.models.database.user_balance import UserBalance +from shared.services.billing.credits_service import CreditsService from shared.services.redis.redis_service import RedisService _DEFAULT_TIER: str = "free" @@ -42,7 +43,13 @@ async def get_tier(user_id: str) -> str: return cached_tier async with get_db_context() as session: - user_tier = await TierService._get_tier_from_db(session, user_id) + try: + user_tier: str = await TierService._get_tier_from_db(session, user_id) + except NotFoundException: + user_tier = await TierService._initialize_missing_user_tier( + session, + user_id, + ) await TierService._set_cached_tier(redis_service, user_id, user_tier) return user_tier @@ -126,6 +133,21 @@ async def _get_tier_from_db(session: AsyncSession, user_id: str) -> str: ) return str(user_tier) + @staticmethod + async def _initialize_missing_user_tier( + session: AsyncSession, + user_id: str, + ) -> str: + """Create missing first-use billing state, then return the user's tier.""" + credits_service: CreditsService = CreditsService() + await credits_service.ensure_user_initialized(session, user_id) + user_tier: str = await TierService._get_tier_from_db(session, user_id) + logger.info( + "Initialized missing user balance during tier lookup: user_id={}", + user_id, + ) + return user_tier + @staticmethod async def _get_cached_tier( redis_service: RedisService, diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 9a9a48469..b62fd941d 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -1,4 +1,5 @@ import importlib +import json from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone @@ -10,12 +11,51 @@ from pytest import MonkeyPatch from tests.support.contract_database import ContractDatabase +from shared.utils.api_keys import hash_api_key def _utc_now() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) +async def _insert_api_key_for_user(user_id: str, api_key: str) -> None: + timestamp = _utc_now() + api_key_id = f"key_{uuid4().hex[:12]}" + await ContractDatabase.execute( + """ + INSERT INTO api_keys ( + id, + user_id, + key_hash, + key_mask, + name, + enabled_modules, + is_active, + created_at + ) VALUES ( + :id, + :user_id, + :key_hash, + :key_mask, + :name, + CAST(:enabled_modules AS JSON), + :is_active, + :created_at + ) + """, + { + "id": api_key_id, + "user_id": user_id, + "key_hash": hash_api_key(api_key), + "key_mask": f"{api_key[:8]}...{api_key[-4:]}", + "name": f"Contract API Key {user_id}", + "enabled_modules": json.dumps(["all"]), + "is_active": True, + "created_at": timestamp, + }, + ) + + @pytest.mark.asyncio async def test_should_return_the_authenticated_users_initialized_credits_balance( developer_api_client_factory: Callable[ @@ -29,6 +69,60 @@ async def test_should_return_the_authenticated_users_initialized_credits_balance assert response.json() == {"credits_balance": 5.0} +@pytest.mark.asyncio +async def test_should_initialize_missing_user_balance_during_tier_lookup( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], +) -> None: + user_id = f"contract-missing-balance-{uuid4().hex[:12]}" + api_key = f"sk_contract_{uuid4().hex[:24]}" + + async with api_client_factory() as api_client: + await ContractDatabase.insert_user(user_id=user_id) + await _insert_api_key_for_user(user_id, api_key) + api_client.headers.update({"Authorization": f"Bearer {api_key}"}) + + response = await api_client.get("/api/v1/billing/credits") + balance_row = await ContractDatabase.fetch_one( + """ + SELECT credits_balance, user_tier + FROM user_balances + WHERE user_id = :user_id + """, + {"user_id": user_id}, + ) + transaction_row = await ContractDatabase.fetch_one( + """ + SELECT credits_amount, transaction_type + FROM credits_transactions + WHERE user_id = :user_id + AND transaction_type = 'initial_grant' + """, + {"user_id": user_id}, + ) + payment_row = await ContractDatabase.fetch_one( + """ + SELECT credits_amount, payment_type, status + FROM payment_records + WHERE user_id = :user_id + AND payment_type = 'system_grant' + """, + {"user_id": user_id}, + ) + + assert response.status_code == 200 + assert response.json() == {"credits_balance": 5.0} + assert balance_row == {"credits_balance": 5_000_000, "user_tier": "free"} + assert transaction_row == { + "credits_amount": 5_000_000, + "transaction_type": "initial_grant", + } + assert payment_row == { + "credits_amount": 5_000_000, + "payment_type": "system_grant", + "status": "succeeded", + } + + @pytest.mark.asyncio async def test_should_not_register_billing_routes_when_billing_is_disabled( monkeypatch: MonkeyPatch, diff --git a/packages/shared-python/shared/services/billing/credits_service.py b/packages/shared-python/shared/services/billing/credits_service.py index 5bf816f84..566440157 100644 --- a/packages/shared-python/shared/services/billing/credits_service.py +++ b/packages/shared-python/shared/services/billing/credits_service.py @@ -48,11 +48,10 @@ class CreditsService: not ensure the /billing/credits endpoint is called before any credit modification operation. - 2. **Read-only endpoints** (e.g., GET /billing/credits): - MUST be called explicitly by the API route. This is necessary because - `get_balance()` is intentionally kept fast (no initialization check) - for performance. Without this call in the route, first-time users - would see 0 balance instead of their initial credits. + 2. **First-use user flows**: + Called before reading balance data, either from an API route or from + the tier lookup path used by authenticated route guards. `get_balance()` + is intentionally kept fast (no initialization check) for performance. Usage: ------