Staging

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,

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,

apps/api/tests/contract/test_documents_contract.py

@@ -105,6 +105,144 @@ async def _fetch_document(document_id: str) -> dict[str, object]:
         await engine.dispose()
 
 
+async def _fetch_graph_counts(
+    *,
+    document_id: str,
+    peer_document_id: str | None = None,
+) -> dict[str, int]:
+    engine = await _create_contract_engine()
+    document_node_id = f"doc:{document_id}"
+    try:
+        async with engine.begin() as connection:
+            node_count = (
+                await connection.execute(
+                    text("""
+                    SELECT COUNT(*)
+                    FROM graph_nodes
+                    WHERE owner_document_id = :document_id
+                    """),
+                    {"document_id": document_id},
+                )
+            ).scalar_one()
+            related_edge_count = (
+                await connection.execute(
+                    text("""
+                    SELECT COUNT(*)
+                    FROM graph_edges
+                    WHERE owner_document_id = :document_id
+                       OR source_node_id = :document_node_id
+                       OR target_node_id = :document_node_id
+                    """),
+                    {
+                        "document_id": document_id,
+                        "document_node_id": document_node_id,
+                    },
+                )
+            ).scalar_one()
+            peer_node_count = 0
+            if peer_document_id is not None:
+                peer_node_count = (
+                    await connection.execute(
+                        text("""
+                        SELECT COUNT(*)
+                        FROM graph_nodes
+                        WHERE owner_document_id = :peer_document_id
+                        """),
+                        {"peer_document_id": peer_document_id},
+                    )
+                ).scalar_one()
+            return {
+                "nodes": int(node_count),
+                "related_edges": int(related_edge_count),
+                "peer_nodes": int(peer_node_count),
+            }
+    finally:
+        await engine.dispose()
+
+
+async def _insert_document_graph_fixture(
+    *,
+    document_id: str,
+    job_result_id: str,
+    peer_document_id: str,
+    peer_job_result_id: str,
+    user_id: str = "local-dev-user",
+    namespace: str = "contract-documents",
+) -> None:
+    engine = await _create_contract_engine()
+    timestamp = datetime.now(timezone.utc).replace(tzinfo=None)
+    try:
+        async with engine.begin() as connection:
+            await connection.execute(
+                text("""
+                INSERT INTO graph_nodes (
+                    node_id,
+                    user_id,
+                    namespace,
+                    node_kind,
+                    owner_document_id,
+                    job_result_id,
+                    ref_document_id,
+                    ref_section_id,
+                    properties,
+                    created_at,
+                    updated_at
+                ) VALUES
+                    (:doc_node_id, :user_id, :namespace, 'document', :document_id, :job_result_id, :document_id, NULL, CAST('{}' AS JSON), :created_at, :updated_at),
+                    (:peer_node_id, :user_id, :namespace, 'document', :peer_document_id, :peer_job_result_id, :peer_document_id, NULL, CAST('{}' AS JSON), :created_at, :updated_at)
+                """),
+                {
+                    "doc_node_id": f"doc:{document_id}",
+                    "peer_node_id": f"doc:{peer_document_id}",
+                    "user_id": user_id,
+                    "namespace": namespace,
+                    "document_id": document_id,
+                    "peer_document_id": peer_document_id,
+                    "job_result_id": job_result_id,
+                    "peer_job_result_id": peer_job_result_id,
+                    "created_at": timestamp,
+                    "updated_at": timestamp,
+                },
+            )
+            await connection.execute(
+                text("""
+                INSERT INTO graph_edges (
+                    edge_id,
+                    user_id,
+                    namespace,
+                    edge_kind,
+                    source_node_id,
+                    target_node_id,
+                    owner_document_id,
+                    job_result_id,
+                    is_directed,
+                    weight,
+                    properties,
+                    created_at,
+                    updated_at
+                ) VALUES
+                    (:owned_edge_id, :user_id, :namespace, 'related', :doc_node_id, :peer_node_id, :document_id, :job_result_id, FALSE, 1.0, CAST('{}' AS JSON), :created_at, :updated_at),
+                    (:incoming_edge_id, :user_id, :namespace, 'related', :peer_node_id, :doc_node_id, :peer_document_id, :peer_job_result_id, FALSE, 1.0, CAST('{}' AS JSON), :created_at, :updated_at)
+                """),
+                {
+                    "owned_edge_id": f"edge_{uuid4().hex[:12]}",
+                    "incoming_edge_id": f"edge_{uuid4().hex[:12]}",
+                    "user_id": user_id,
+                    "namespace": namespace,
+                    "doc_node_id": f"doc:{document_id}",
+                    "peer_node_id": f"doc:{peer_document_id}",
+                    "document_id": document_id,
+                    "peer_document_id": peer_document_id,
+                    "job_result_id": job_result_id,
+                    "peer_job_result_id": peer_job_result_id,
+                    "created_at": timestamp,
+                    "updated_at": timestamp,
+                },
+            )
+    finally:
+        await engine.dispose()
+
+
 async def _insert_document_revision_with_chunks(
     *,
     document_id: str,
@@ -662,21 +800,62 @@ async def test_should_archive_a_document_via_the_canonical_archive_route(
     ],
 ) -> None:
     document_id = f"doc_{uuid4().hex[:12]}"
+    peer_document_id = f"doc_{uuid4().hex[:12]}"
 
     async with developer_api_client_factory() as api_client:
-        await _insert_document(document_id=document_id)
+        document_revision = await _insert_document_revision_with_chunks(
+            document_id=document_id,
+            chunks=[
+                {
+                    "id": f"dchk_{uuid4().hex[:12]}",
+                    "chunk_id": "archive-chunk-1",
+                    "chunk_type": "text",
+                    "content": "Archived graph chunk",
+                    "source_chunk_path": "Chapter 1/Archive",
+                    "metadata": {"keywords": ["archive"]},
+                }
+            ],
+        )
+        peer_revision = await _insert_document_revision_with_chunks(
+            document_id=peer_document_id,
+            chunks=[
+                {
+                    "id": f"dchk_{uuid4().hex[:12]}",
+                    "chunk_id": "peer-chunk-1",
+                    "chunk_type": "text",
+                    "content": "Peer graph chunk",
+                    "source_chunk_path": "Chapter 1/Peer",
+                    "metadata": {"keywords": ["peer"]},
+                }
+            ],
+        )
+        await _insert_document_graph_fixture(
+            document_id=document_id,
+            job_result_id=document_revision["job_result_id"],
+            peer_document_id=peer_document_id,
+            peer_job_result_id=peer_revision["job_result_id"],
+        )
         response = await api_client.post(f"/api/v1/documents/{document_id}/archive")
 
     assert response.status_code == 200
 
     response_json = cast(dict[str, object], response.json())
     persisted_document = await _fetch_document(document_id)
+    graph_counts = await _fetch_graph_counts(
+        document_id=document_id,
+        peer_document_id=peer_document_id,
+    )
 
     assert response_json["document_id"] == document_id
     assert response_json["status"] == "archived"
     assert response_json["archived_at"]
     assert persisted_document["status"] == "archived"
     assert persisted_document["archived_at"] is not None
+    assert graph_counts == {
+        "nodes": 0,
+        "related_edges": 0,
+        "peer_nodes": 1,
+    }
 
 
 @pytest.mark.asyncio