M2: atomic reserve-before-work quota check

app/routes/v1.py

@@ -148,8 +148,12 @@ async def _handle_billable_literature_call(
             return _error(409, "idempotency_in_progress",
                           "A request with this Idempotency-Key is still running.", request_id, headers)
 
-    remaining = await store.quota_remaining(org_id, resource)
-    if remaining is not None and remaining <= 0:
+    # Atomic reserve-before-work: decrements the quota counter now, not
+    # after the upstream call succeeds, so concurrent requests can never
+    # all observe "quota available" and all succeed (see
+    # V1Store.reserve_quota's own docstring for why the old check-then-
+    # later-decrement pair could overrun a near-zero quota).
+    if not await store.reserve_quota(org_id, resource):
         if idempotency_key is not None:
             await store.idempotency_finish(subject, idempotency_key, fingerprint, 402, None)
         return _error(402, "quota_exceeded", quota_exceeded_message, request_id, headers)
@@ -159,6 +163,9 @@ async def _handle_billable_literature_call(
     latency_ms = round((time.monotonic() - started) * 1000)
 
     if not 200 <= status < 300:
+        # The reservation above assumed this call would succeed and be
+        # billed; it didn't, so the unit must be given back.
+        await store.release_quota(org_id, resource)
         if idempotency_key is not None:
             await store.idempotency_finish(subject, idempotency_key, fingerprint, status, response)
         return _upstream_error(status, response, request_id, headers)
@@ -183,7 +190,6 @@ async def _handle_billable_literature_call(
             "idempotency_key_sha256": request_fingerprint(idempotency_key) if idempotency_key else None,
         },
     )
-    await store.consume_quota(org_id, resource)
     return JSONResponse(public_response, status_code=status, headers={**headers, "X-Request-Id": request_id})
 
 

app/services/v1_store.py

@@ -140,13 +140,52 @@ async def quota_remaining(self, org_id: str, resource: str) -> Optional[int]:
         except ValueError:
             return None
 
-    async def consume_quota(self, org_id: str, resource: str) -> None:
-        """Decrement an existing quota after a billed success. Never creates
-        the key: an org without a quota stays unmetered here."""
+    async def reserve_quota(self, org_id: str, resource: str) -> bool:
+        """Atomically reserve one unit of `resource` for `org_id` before
+        doing the work it would bill for, and report whether the
+        reservation succeeded. Replaces the old quota_remaining()-then-
+        later-consume_quota() pair, which read the counter, did the
+        (slow) upstream call, and only decremented afterward -- any
+        number of concurrent requests could all observe "1 remaining"
+        before any of them decremented, and all of them would then
+        succeed, overrunning the quota by however many were in flight
+        at once. Redis's DECR is atomic, so only as many concurrent
+        reservations as there are units left can ever observe a
+        non-negative result here; the rest observe negative and
+        compensate back to zero immediately (never below zero, and
+        never creating a key that did not already exist -- an org
+        without a quota key stays unmetered, exactly like before).
+
+        Fails open (reservation succeeds) on a Redis error, the same
+        posture every other method in this class takes: an outage must
+        never block a paid request, only leave this particular overrun
+        protection briefly unenforced until Redis recovers.
+        """
+        key = self._quota_key(org_id, resource)
+        try:
+            if await self.redis.get(key) is None:
+                return True  # unmetered: no quota key set for this org/resource
+            value = await self.redis.decr(key)
+        except Exception:
+            return True
+        if value < 0:
+            try:
+                await self.redis.incr(key)
+            except Exception:
+                pass
+            return False
+        return True
+
+    async def release_quota(self, org_id: str, resource: str) -> None:
+        """Refund a reservation made by reserve_quota() for a call that was
+        then not actually billable (the upstream request failed) -- the
+        reservation already decremented optimistically, before knowing
+        whether the call would succeed. Never creates the key: mirrors
+        reserve_quota's own "only adjust an existing counter" rule."""
         key = self._quota_key(org_id, resource)
         try:
             if await self.redis.get(key) is not None:
-                await self.redis.decr(key)
+                await self.redis.incr(key)
         except Exception:
             pass
 

tests/test_v1_literature.py

@@ -242,6 +242,38 @@ def test_quota_is_consumed_on_success_and_never_created(client, redis, upstream)
     assert "gateway:v1:quota:42:literature.answer" not in redis.kv
 
 
+def test_quota_reservation_is_atomic_under_concurrency(client, redis, upstream):
+    """Two concurrent requests with only 1 unit of quota left: exactly one
+    must succeed and one must be quota_exceeded -- not both succeeding
+    (an overrun) and not both failing (undercounting real capacity).
+    FakeRedis's incr/decr aren't async-concurrent in the true sense (no
+    real parallelism in this test process), but this still exercises the
+    actual reserve-then-compensate sequence reserve_quota runs, not a
+    mock standing in for it."""
+    redis.kv["gateway:v1:quota:42:literature.answer"] = 1
+    first = _post(client, idem="race-1")
+    second = _post(client, idem="race-2")
+    statuses = sorted([first.status_code, second.status_code])
+    assert statuses == [200, 402]
+    # The quota key never goes negative and ends at exactly zero, not
+    # some other value a non-atomic check-then-decrement could leave it at.
+    assert redis.kv["gateway:v1:quota:42:literature.answer"] == 0
+
+
+def test_quota_is_refunded_when_upstream_call_fails(client, redis, upstream):
+    """reserve_quota decrements optimistically, before knowing whether the
+    call will succeed -- a failed upstream call must give the unit back,
+    or a string of transient RAG failures would silently burn through an
+    organization's quota for answers it never actually got billed for
+    (and never received)."""
+    redis.kv["gateway:v1:quota:42:literature.answer"] = 1
+    upstream.return_value = (503, {"detail": "rag down"})
+    resp = _post(client)
+    assert resp.status_code == 503
+    assert redis.kv["gateway:v1:quota:42:literature.answer"] == 1
+    assert _usage(redis) == []
+
+
 def test_idempotent_retry_replays_without_running_or_billing_again(client, redis, upstream):
     first = _post(client, idem="retry-1")
     second = _post(client, idem="retry-1")
@@ -362,7 +394,7 @@ def test_search_rejects_missing_question(client, redis, upstream):
 
 
 def test_search_quota_exceeded(client, redis, upstream):
-    with patch.object(v1.store, "quota_remaining", AsyncMock(return_value=0)):
+    with patch.object(v1.store, "reserve_quota", AsyncMock(return_value=False)):
         resp = client.post("/v1/literature/search", json={"question": "q"},
                             headers={"Authorization": f"Bearer {KEY}"})
     assert resp.status_code == 402