feat: token-bucket rate limiting and concurrent-answer limits

app/core/config.py

@@ -42,6 +42,16 @@ class Config:
     V1_REDIS_URL = os.getenv("V1_REDIS_URL", "")
     V1_RATE_LIMIT_PER_MINUTE = int(os.getenv("V1_RATE_LIMIT_PER_MINUTE", "60"))
     V1_IDEMPOTENCY_TTL = int(os.getenv("V1_IDEMPOTENCY_TTL", "86400"))
+    # Design audit gap #7 ("concurrent-answer limits are absent"): the
+    # most expensive /v1 call (it invokes an LLM, up to RAG's own
+    # 300-second timeout) is capped on how many of a single key's or
+    # organization's calls may be in flight at once -- independent of,
+    # and enforced alongside, the per-minute rate limit above, which
+    # only bounds call *frequency*, not concurrency. A conservative
+    # operator-tunable default, not a per-plan number (no billing_plans
+    # column exists for this, unlike V1_RATE_LIMIT_PER_MINUTE's own
+    # plan-aware override -- this is overload protection, not pricing).
+    V1_MAX_CONCURRENT_ANSWERS = int(os.getenv("V1_MAX_CONCURRENT_ANSWERS", "5"))
     # Billable usage goes to the same usage:events stream every other
     # producer writes (omnibioai-usage-client wire format), consumed by
     # omnibioai-billing.

app/routes/v1.py

@@ -5,9 +5,14 @@
 against. On top of the middleware chain (auth incl. omni_sk_ API keys,
 policy, audit) every /v1 call gets:
 
-- a rate limit enforced both per caller and per organization (X-RateLimit-*
-  headers, 429 + Retry-After) -- an organization can't multiply its
-  effective limit by spreading requests across several API keys,
+- a token-bucket rate limit enforced both per caller and per organization
+  (X-RateLimit-* headers, 429 + Retry-After) -- an organization can't
+  multiply its effective limit by spreading requests across several API
+  keys, and bursting a full minute's allowance at once no longer lets a
+  caller squeeze in double that across one window boundary,
+- a concurrency limit on in-flight /v1/literature/answers calls, per
+  caller and per organization (429, independent of the rate limit above,
+  which only bounds call frequency),
 - an org-level quota check maintained by omnibioai-billing (402),
 - optional Idempotency-Key replay, so a retried request is never run or
   billed twice,
@@ -109,6 +114,26 @@ async def _rate_limited(request: Request, subject: str, org_id: str, request_id:
     return headers, None
 
 
+async def _acquire_concurrency_slots(org_id: str, subject: str, limit: int) -> bool:
+    """Both the subject's and the organization's in-flight counters must
+    have a free slot for the duration of one /v1/literature/answers call
+    -- if the organization's is full, a key that's never made a request
+    of its own must still be blocked, the same per-key-and-per-org
+    pairing _rate_limited already enforces for request frequency."""
+    if not await store.acquire_concurrency_slot(subject, limit):
+        return False
+    if org_id and not await store.acquire_concurrency_slot(f"org:{org_id}", limit):
+        await store.release_concurrency_slot(subject)
+        return False
+    return True
+
+
+async def _release_concurrency_slots(org_id: str, subject: str) -> None:
+    await store.release_concurrency_slot(subject)
+    if org_id:
+        await store.release_concurrency_slot(f"org:{org_id}")
+
+
 async def _forward_to(service: str, request: Request, method: str, path: str, body=None):
     url = f"{resolve_service(service)}/{path}"
     if request.url.query:
@@ -130,6 +155,7 @@ def _upstream_error(status: int, response, request_id: str, headers: dict):
 
 async def _handle_billable_literature_call(
     request: Request, *, resource: str, build_rag_body, build_public_response, quota_exceeded_message: str,
+    max_concurrent_answers: int | None = None,
 ):
     """Shared lifecycle for every billable /v1/literature/* call
     (currently /answers and /search): auth'd-org check, contract
@@ -180,49 +206,67 @@ async def _handle_billable_literature_call(
             return _error(409, "idempotency_in_progress",
                           "A request with this Idempotency-Key is still running.", request_id, headers)
 
-    # 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)
-
-    started = time.monotonic()
-    status, response = await _forward(request, "POST", "v1/query", rag_body)
-    latency_ms = round((time.monotonic() - started) * 1000)
+    if max_concurrent_answers is not None:
+        if not await _acquire_concurrency_slots(org_id, subject, max_concurrent_answers):
+            if idempotency_key is not None:
+                await store.idempotency_finish(subject, idempotency_key, fingerprint, 429, None)
+            return _error(
+                429, "concurrency_limit_exceeded",
+                f"More than {max_concurrent_answers} concurrent requests for this key or organization.",
+                request_id, headers,
+            )
 
-    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)
+    try:
+        # 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)
+
+        started = time.monotonic()
+        status, response = await _forward(request, "POST", "v1/query", rag_body)
+        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)
+
+        public_response = build_public_response(
+            response, domain=body.get("domain"), request_id=request_id, latency_ms=latency_ms,
+        )
         if idempotency_key is not None:
-            await store.idempotency_finish(subject, idempotency_key, fingerprint, status, response)
-        return _upstream_error(status, response, request_id, headers)
-
-    public_response = build_public_response(
-        response, domain=body.get("domain"), request_id=request_id, latency_ms=latency_ms,
-    )
-    if idempotency_key is not None:
-        await store.idempotency_finish(subject, idempotency_key, fingerprint, status, public_response)
-
-    identity = getattr(request.state, "identity", None) or {}
-    await store.emit_usage(
-        org_id=org_id,
-        user_id=user_id,
-        resource=resource,
-        trace_id=request_id,
-        dedup_key=f"{subject}:{idempotency_key}" if idempotency_key else request_id,
-        metadata={
-            "request_id": request_id,
-            "client_id": identity.get("client_id"),
-            "token_type": identity.get("token_type"),
-            "idempotency_key_sha256": request_fingerprint(idempotency_key) if idempotency_key else None,
-        },
-    )
-    return JSONResponse(public_response, status_code=status, headers={**headers, "X-Request-Id": request_id})
+            await store.idempotency_finish(subject, idempotency_key, fingerprint, status, public_response)
+
+        identity = getattr(request.state, "identity", None) or {}
+        await store.emit_usage(
+            org_id=org_id,
+            user_id=user_id,
+            resource=resource,
+            trace_id=request_id,
+            dedup_key=f"{subject}:{idempotency_key}" if idempotency_key else request_id,
+            metadata={
+                "request_id": request_id,
+                "client_id": identity.get("client_id"),
+                "token_type": identity.get("token_type"),
+                "idempotency_key_sha256": request_fingerprint(idempotency_key) if idempotency_key else None,
+            },
+        )
+        return JSONResponse(public_response, status_code=status, headers={**headers, "X-Request-Id": request_id})
+    finally:
+        # Released regardless of how the try block above exited (a quota
+        # denial, an upstream error, or success) -- a slot held by a
+        # request that's already finished answering must never count
+        # against the next one.
+        if max_concurrent_answers is not None:
+            await _release_concurrency_slots(org_id, subject)
 
 
 @router.post("/literature/answers")
@@ -234,6 +278,7 @@ async def literature_answers(request: Request):
         build_public_response=build_public_answer,
         quota_exceeded_message="Your organization has used its included answers. "
                                 "Add a payment method or upgrade the plan.",
+        max_concurrent_answers=Config.V1_MAX_CONCURRENT_ANSWERS,
     )
 
 

app/services/v1_store.py

@@ -6,6 +6,7 @@
 """
 import hashlib
 import json
+import math
 import time
 import uuid
 from datetime import datetime, timezone
@@ -106,21 +107,94 @@ def __init__(self, redis_url: str, usage_redis_url: str, outbox_path: str = Conf
         self.usage_redis = aioredis.from_url(usage_redis_url, decode_responses=True)
         self.outbox = Outbox(outbox_path)
 
-    # ---------------- rate limit (fixed one-minute window) ----------------
+    # ---------------- rate limit (token bucket) ----------------------------
     async def hit_rate_limit(self, subject: str, limit: int) -> tuple[bool, int, int]:
-        """Count one request for `subject`. Returns (allowed, remaining,
-        seconds_until_reset). Fails open on Redis errors."""
-        now = int(time.time())
-        window = now // 60
-        reset = 60 - (now % 60)
-        key = f"{_PREFIX}rl:{subject}:{window}"
+        """Design audit gap #7 ("fixed one-minute window... rather than a
+        token bucket"): `limit` tokens refill continuously over 60 seconds
+        (rate = limit/60 tokens/sec), bucket capacity = limit -- a caller
+        can burst its full per-minute allowance at once, then it refills
+        smoothly, rather than the old fixed window's hard reset-at-:00
+        boundary, which let a caller spend its whole budget in the last
+        second of one window and again in the first second of the next:
+        up to 2x `limit` requests in under two seconds, never caught by a
+        window that only ever compares against *one* window's count at a
+        time. Returns (allowed, remaining, seconds_until_next_token) --
+        remaining is the floored token count after this request; the third
+        value is 0 whenever a request is allowed (another token is already
+        available right now) and otherwise how long until the bucket has
+        earned back at least one, used for X-RateLimit-Reset/Retry-After.
+
+        Not atomic against a concurrent request for the *same* subject on
+        a real Redis backend (a plain GET then SET, no Lua script) -- a
+        request landing in that window could read the same token count and
+        both deduct from it, each believing it got a distinct token. A
+        bounded, accepted imperfection (the practical cost is occasionally
+        allowing one extra request under true concurrent load for one
+        subject), not a new class of risk: nothing here is billing-
+        critical the way reserve_quota's own atomic DECR had to be. Fails
+        open on a Redis error, same posture as every other method here.
+        """
+        if limit <= 0:
+            return False, 0, 60
+        rate = limit / 60.0
+        now = time.time()
+        key = f"{_PREFIX}rltb:{subject}"
+        try:
+            raw = await self.redis.get(key)
+            if raw is None:
+                tokens, last = float(limit), now
+            else:
+                tokens, last = json.loads(raw)
+            tokens = min(float(limit), tokens + max(0.0, now - last) * rate)
+            allowed = tokens >= 1.0
+            if allowed:
+                tokens -= 1.0
+            await self.redis.set(key, json.dumps([tokens, now]), ex=120)
+        except Exception:
+            return True, limit, 0
+        reset = 0 if tokens >= 1.0 else max(1, math.ceil((1.0 - tokens) / rate))
+        return allowed, int(tokens), reset
+
+    # ---------------- concurrent-request limit ------------------------------
+    def _concurrency_key(self, subject: str) -> str:
+        return f"{_PREFIX}conc:{subject}"
+
+    async def acquire_concurrency_slot(self, subject: str, limit: int) -> bool:
+        """Design audit gap #7 ("concurrent-answer limits are absent"):
+        atomically increments the in-flight-request counter for `subject`
+        (an API key, a user session, or f"org:{org_id}") and reports
+        whether the new count is within `limit`. If not, immediately
+        decrements back out -- a caller denied a slot never holds one, so
+        it must call release_concurrency_slot only when this returns True
+        (the same reserve-then-release pairing reserve_quota/release_quota
+        already established for the billing quota counter). The TTL is a
+        safety net only, for the case a crash skips the matching release --
+        every normal request releases its own slot long before 300s.
+        Fails open (True) on a Redis error.
+        """
+        key = self._concurrency_key(subject)
         try:
             count = await self.redis.incr(key)
-            if count == 1:
-                await self.redis.expire(key, 61)
+            await self.redis.expire(key, 300)
+        except Exception:
+            return True
+        if count > limit:
+            try:
+                await self.redis.decr(key)
+            except Exception:
+                pass
+            return False
+        return True
+
+    async def release_concurrency_slot(self, subject: str) -> None:
+        """Releases a slot acquired by acquire_concurrency_slot. Must only
+        be called for a subject that call actually returned True for --
+        calling it for a denied acquisition would double-release (the
+        denial already decremented back out itself)."""
+        try:
+            await self.redis.decr(self._concurrency_key(subject))
         except Exception:
-            return True, limit, reset
-        return count <= limit, max(0, limit - count), reset
+            pass
 
     # ---------------- rate limit override (maintained by omnibioai-billing)
     def _org_rate_limit_key(self, org_id: str) -> str: