fix: backfill empty entitlement account_id from DCR and retry on 5xx

src/lightspeed_agent/dcr/service.py

@@ -172,6 +172,11 @@ async def register_client(
                 error_description=f"Invalid Order ID: {claims.order_id}",
             )
 
+        # Step 3b: Backfill entitlement account_id if empty
+        await self._procurement_service.backfill_entitlement_account_id(
+            claims.order_id, claims.account_id
+        )
+
         # Step 4: Check if client already exists for this order
         existing_client = await self._client_repository.get_by_order_id(claims.order_id)
         if existing_client:

src/lightspeed_agent/marketplace/service.py

@@ -459,6 +459,37 @@ async def _handle_offer_ended(self, event: ProcurementEvent) -> None:
         if event.entitlement:
             logger.info("Offer ended: %s", event.entitlement.id)
 
+    # Public helpers
+
+    async def backfill_entitlement_account_id(self, order_id: str, account_id: str) -> None:
+        """Backfill an entitlement's account_id if it is currently empty.
+
+        Called during DCR registration when the JWT provides an authoritative
+        account_id that may be missing from the entitlement record (Pub/Sub
+        events often arrive without account info).
+
+        Args:
+            order_id: The entitlement/order ID.
+            account_id: The account ID to set (from the DCR JWT ``sub`` claim).
+        """
+        if not account_id:
+            return
+
+        entitlement = await self._entitlement_repo.get(order_id)
+        if not entitlement:
+            return
+
+        if entitlement.account_id:
+            return
+
+        entitlement.account_id = account_id
+        await self._entitlement_repo.update(entitlement)
+        logger.info(
+            "Backfilled account_id %s on entitlement %s from DCR",
+            account_id,
+            order_id,
+        )
+
     # Procurement API operations
 
     async def _resolve_account_id(self, entitlement_id: str, event: ProcurementEvent) -> str | None:
@@ -519,6 +550,11 @@ async def _resolve_account_id(self, entitlement_id: str, event: ProcurementEvent
                 "Entitlement %s not found in Procurement API (404)",
                 entitlement_id,
             )
+        elif response.status_code >= 500:
+            raise RuntimeError(
+                f"Server error resolving account for entitlement {entitlement_id}: "
+                f"HTTP {response.status_code} — {response.text}"
+            )
         else:
             logger.warning(
                 "Failed to fetch entitlement %s: HTTP %s — %s",

tests/test_dcr.py

@@ -172,6 +172,74 @@ async def test_get_client(self, service):
         assert client.account_id == "valid-account-123"
 
 
+class TestDCRServiceBackfill:
+    """Tests for DCR service backfilling entitlement account_id."""
+
+    @pytest_asyncio.fixture
+    async def service_and_repo(self, db_session):
+        """Create a DCR service with an entitlement that has empty account_id."""
+        entitlement_repo = EntitlementRepository()
+        client_repo = DCRClientRepository()
+        procurement_service = ProcurementService(
+            entitlement_repo=entitlement_repo,
+        )
+
+        entitlement = Entitlement(
+            id="backfill-order-789",
+            account_id="",
+            provider_id="provider-456",
+            state=EntitlementState.ACTIVE,
+        )
+        await entitlement_repo.create(entitlement)
+
+        svc = DCRService(
+            procurement_service=procurement_service,
+            client_repository=client_repo,
+        )
+        return svc, entitlement_repo
+
+    @pytest.mark.asyncio
+    async def test_register_client_backfills_entitlement_account_id(self, service_and_repo):
+        """Test that DCR registration backfills empty entitlement account_id."""
+        from lightspeed_agent.dcr.gma_client import GMAClientResponse
+
+        service, entitlement_repo = service_and_repo
+
+        claims = GoogleJWTClaims(
+            iss="https://accounts.google.com",
+            iat=int(time.time()),
+            exp=int(time.time()) + 3600,
+            aud="https://example.com",
+            sub="backfill-account-123",
+            google={"order": "backfill-order-789"},
+            auth_app_redirect_uris=["https://example.com/callback"],
+        )
+
+        mock_gma = AsyncMock()
+        mock_gma.create_tenant.return_value = GMAClientResponse(
+            client_id="new-client-id",
+            client_secret="new-secret",
+            name="lightspeed-backfill-order-789",
+        )
+        service._gma_client = mock_gma
+
+        # Mock JWT validation to return our claims
+        with patch.object(
+            service._jwt_validator,
+            "validate_software_statement",
+            new_callable=AsyncMock,
+            return_value=claims,
+        ):
+            result = await service.register_client(
+                DCRRequest(software_statement="valid-jwt")
+            )
+
+        assert isinstance(result, DCRResponse)
+
+        ent = await entitlement_repo.get("backfill-order-789")
+        assert ent.account_id == "backfill-account-123"
+
+
 class TestDCRServiceEncryptionValidation:
     """Tests for DCR service encryption key validation."""
 

tests/test_marketplace.py

@@ -373,8 +373,72 @@ async def test_entitlement_creation_idempotent(self, service):
             mock_settings.skip_pubsub_oidc_verification = False
             await service.process_event(event)
 
+    # --- backfill_entitlement_account_id tests ---
+
+    @pytest.mark.asyncio
+    async def test_backfill_entitlement_account_id(self, service):
+        """Test backfill updates empty account_id."""
+        ent = Entitlement(
+            id="order-backfill",
+            account_id="",
+            state=EntitlementState.ACTIVE,
+            provider_id="provider-1",
+        )
+        await service._entitlement_repo.create(ent)
+
+        await service.backfill_entitlement_account_id("order-backfill", "account-filled")
+
+        updated = await service._entitlement_repo.get("order-backfill")
+        assert updated.account_id == "account-filled"
+
+    @pytest.mark.asyncio
+    async def test_backfill_entitlement_account_id_no_overwrite(self, service):
+        """Test backfill does not overwrite existing account_id."""
+        ent = Entitlement(
+            id="order-existing-acct",
+            account_id="original-account",
+            state=EntitlementState.ACTIVE,
+            provider_id="provider-1",
+        )
+        await service._entitlement_repo.create(ent)
+
+        await service.backfill_entitlement_account_id("order-existing-acct", "new-account")
+
+        updated = await service._entitlement_repo.get("order-existing-acct")
+        assert updated.account_id == "original-account"
+
+    @pytest.mark.asyncio
+    async def test_backfill_entitlement_account_id_not_found(self, service):
+        """Test backfill on non-existent order does not raise."""
+        await service.backfill_entitlement_account_id("order-nonexistent", "account-1")
+
     # --- _resolve_account_id tests ---
 
+    @pytest.mark.asyncio
+    async def test_resolve_account_id_raises_on_5xx(self, service):
+        """Test _resolve_account_id raises RuntimeError on 5xx server errors."""
+        event = ProcurementEvent(
+            event_id="event-1",
+            event_type=ProcurementEventType.ENTITLEMENT_CREATION_REQUESTED,
+            provider_id="provider-1",
+            entitlement=EntitlementInfo(id="order-1"),
+        )
+
+        mock_response = httpx.Response(
+            status_code=500,
+            text="Internal Server Error",
+            request=httpx.Request("GET", "https://example.com"),
+        )
+        with (
+            patch.object(service, "_settings") as mock_settings,
+            patch.object(service, "_get_auth_headers", return_value={}),
+            patch("httpx.AsyncClient.get", new_callable=AsyncMock, return_value=mock_response),
+        ):
+            mock_settings.google_cloud_project = "test-project"
+            mock_settings.skip_pubsub_oidc_verification = False
+            with pytest.raises(RuntimeError, match="Server error resolving account"):
+                await service._resolve_account_id("order-1", event)
+
     @pytest.mark.asyncio
     async def test_resolve_account_id_from_event(self, service):
         """Test _resolve_account_id returns account ID from event payload."""