fix(interactions): keep the stored settlement gate when registration raises after landing, and fail resumed-poll deletes closed

A registration that raised after its row committed moved the poll to a private in-memory gate, so the creating worker billed while the stored row stayed unclaimed for another replica's delete or the next boot to bill again. The row is now read back once and, when it landed, the poll claims through it like every other settler.

A worker that resumed the poll after a restart is not the creator, so its delete on a failed pre-delete fetch now fails with the fetch's error instead of releasing and deleting. After a fleet restart every worker holds resumed polls, which left the fail-closed path applying nowhere.
This commit is contained in:
mateo-berri 2026-09-24 13:13:40 -07:00
parent 8b451c6cf9
commit 0a31f7b705
2 changed files with 107 additions and 1 deletions

View file

@ -174,6 +174,7 @@ class BackgroundInteractionPollContext:
max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS
timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS
store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore)
resumed: bool = False
FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Awaitable[InteractionsAPIResponse]]
@ -490,6 +491,12 @@ async def _registered_store(
try:
await store.register(pending)
except Exception: # noqa: BLE001 # a store outage must not fail the create; the poll settles from this process
if await _registration_landed(store, pending.interaction_id):
verbose_logger.exception(
"Registering background interaction %s raised although its row landed; it settles through the store",
pending.interaction_id,
)
return store
verbose_logger.exception(
"Could not durably register background interaction %s; only this process can settle it",
pending.interaction_id,
@ -500,6 +507,13 @@ async def _registered_store(
return store
async def _registration_landed(store: BackgroundSettlementStore, interaction_id: str) -> bool:
try:
return await store.pending(interaction_id) is not None or await store.is_claimed(interaction_id)
except Exception: # noqa: BLE001 # the row cannot be read either, so the poll settles from this process
return False
async def maybe_schedule_background_interaction_cost_polling(
response: object,
create_kwargs: Mapping[str, object],
@ -575,7 +589,7 @@ async def maybe_settle_background_interaction_before_delete(
store: BackgroundSettlementStore | None = None,
) -> SettlementOutcome | None:
entry: Final = _ACTIVE_POLLS.get(interaction_id)
if entry is not None:
if entry is not None and not entry.context.resumed:
return await _settle_before_delete(entry.context, await _fetch_before_delete(entry.context, fetch_interaction))
settlement_store: Final = store or _STORE.store
pending: Final = await _pending(settlement_store, interaction_id)
@ -633,6 +647,7 @@ def _resumed_context(
max_interval_seconds=schedule.max_interval_seconds,
timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds),
store=store,
resumed=True,
)

View file

@ -848,6 +848,97 @@ class _DownStore:
raise RuntimeError("database unavailable")
class _RegistersThenRaises:
def __init__(self):
self.store = InMemoryBackgroundSettlementStore()
async def register(self, pending):
await self.store.register(pending)
raise RuntimeError("connection reset after the row was committed")
async def pending(self, interaction_id):
return await self.store.pending(interaction_id)
async def is_claimed(self, interaction_id):
return await self.store.is_claimed(interaction_id)
async def claim(self, interaction_id):
return await self.store.claim(interaction_id)
async def record_outcome(self, interaction_id, outcome):
return None
async def unclaimed(self):
return await self.store.unclaimed()
@pytest.mark.asyncio
async def test_create_whose_registration_raised_after_landing_still_claims_the_stored_row():
"""
A registration that raises after its row committed used to move the poll
to a private in-memory gate, so the creating worker billed while the
stored row stayed unclaimed for another replica's delete or the next boot
to bill again. The row that landed is the gate every settler shares.
"""
store = _RegistersThenRaises()
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()})
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
task = await maybe_schedule_background_interaction_cost_polling(
response=_response("in_progress", with_usage=False),
create_kwargs={"litellm_logging_obj": logging_obj},
custom_llm_provider="gemini",
store=store,
fetch_interaction=poll_fetch,
)
fetch, _ = _capturing_fetch(_response("completed", with_usage=True))
outcome = await maybe_settle_background_interaction_before_delete(
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
)
assert outcome == "billed"
assert await store.is_claimed("interactions/bg-abc")
assert await store.unclaimed() == ()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
@pytest.mark.asyncio
async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch():
"""
After a restart every worker resumes the unclaimed rows, so none of them
is the creator whose delete may release and delete on a failed fetch. A
resumed worker's delete fails like any other replica's, and its own poll
still bills the interaction once it completes.
"""
store = InMemoryBackgroundSettlementStore()
await store.register(
PendingBackgroundInteraction(
interaction_id="interactions/bg-abc",
custom_llm_provider="gemini",
create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"),
created_at=datetime.now(timezone.utc),
)
)
responses = [_response("in_progress", with_usage=False)]
async def poll_fetch(context):
return responses[-1]
(resumed,) = await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE)
fetch, _ = _fetch_sequence(RuntimeError("Google API key is required"))
with pytest.raises(RuntimeError, match="Google API key is required"):
await maybe_settle_background_interaction_before_delete(
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
)
assert await store.is_claimed("interactions/bg-abc") is False
responses.append(_response("completed", with_usage=True))
assert await asyncio.wait_for(resumed, timeout=5) == "billed"
@pytest.mark.asyncio
async def test_create_still_settles_on_its_own_replica_when_the_store_is_down():
logging_obj = _logging_obj()