From e631e7b615548e1d51deffcf7597a5ca2cf9a7fb Mon Sep 17 00:00:00 2001 From: Luis Tomas Bolivar Date: Wed, 8 Jul 2026 10:12:39 +0200 Subject: [PATCH 1/2] fix: raise on 5xx in _resolve_account_id for Pub/Sub retry Previously, 5xx responses from the Procurement API when resolving account IDs were silently swallowed (logged as warning, returned None). This meant the entitlement was created with an empty account_id and Pub/Sub considered the event handled (200 response), never retrying. Now 5xx responses raise RuntimeError (same pattern as network errors), causing the handler to fail and Pub/Sub to retry the event. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/lightspeed_agent/marketplace/service.py | 5 +++++ tests/test_marketplace.py | 25 +++++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/src/lightspeed_agent/marketplace/service.py b/src/lightspeed_agent/marketplace/service.py index 680c927e..789347fb 100644 --- a/src/lightspeed_agent/marketplace/service.py +++ b/src/lightspeed_agent/marketplace/service.py @@ -519,6 +519,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", diff --git a/tests/test_marketplace.py b/tests/test_marketplace.py index 09c637b0..df09ebdb 100644 --- a/tests/test_marketplace.py +++ b/tests/test_marketplace.py @@ -375,6 +375,31 @@ async def test_entitlement_creation_idempotent(self, service): # --- _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.""" From d015678875347280d40bc8dc300e22b883deaadf Mon Sep 17 00:00:00 2001 From: Luis Tomas Bolivar Date: Wed, 8 Jul 2026 10:50:13 +0200 Subject: [PATCH 2/2] fix: backfill entitlement account_id from DCR registration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dcr_clients table always has a valid account_id (from the JWT sub claim), but marketplace_entitlements often has account_id='' because Pub/Sub events rarely include account info. Add backfill_entitlement_account_id() to ProcurementService and call it during DCR register_client() — every DCR request now writes the JWT's authoritative account_id back to the entitlement if it was previously empty. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/lightspeed_agent/dcr/service.py | 5 ++ src/lightspeed_agent/marketplace/service.py | 31 ++++++++++ tests/test_dcr.py | 68 +++++++++++++++++++++ tests/test_marketplace.py | 39 ++++++++++++ 4 files changed, 143 insertions(+) diff --git a/src/lightspeed_agent/dcr/service.py b/src/lightspeed_agent/dcr/service.py index a28f011b..913b9068 100644 --- a/src/lightspeed_agent/dcr/service.py +++ b/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: diff --git a/src/lightspeed_agent/marketplace/service.py b/src/lightspeed_agent/marketplace/service.py index 789347fb..c2fafeae 100644 --- a/src/lightspeed_agent/marketplace/service.py +++ b/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: diff --git a/tests/test_dcr.py b/tests/test_dcr.py index 21ad31f8..3e874c0f 100644 --- a/tests/test_dcr.py +++ b/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.""" diff --git a/tests/test_marketplace.py b/tests/test_marketplace.py index df09ebdb..229ad623 100644 --- a/tests/test_marketplace.py +++ b/tests/test_marketplace.py @@ -373,6 +373,45 @@ 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