From 0bf39b4b2a49c2a144d57692ccb6954bb3d8a2c9 Mon Sep 17 00:00:00 2001 From: man4ish Date: Sat, 3 Oct 2026 20:18:24 -0500 Subject: [PATCH] feat: token-bucket rate limiting and concurrent-answer limits Replaces the fixed one-minute rate-limit window with a token bucket (continuous refill, same capacity as the configured/plan-aware limit) -- closes the old window's double-burst-across-a-boundary gap. Also adds a concurrency cap on in-flight /v1/literature/answers calls, enforced per key and per organization the same way the rate limiter already is, independent of (and alongside) request-frequency limiting. Both close the two remaining bullets of design audit gap #7. Co-Authored-By: Claude Sonnet 5 --- app/core/config.py | 10 +++ app/routes/v1.py | 133 ++++++++++++++++++++++++------------ app/services/v1_store.py | 96 +++++++++++++++++++++++--- tests/test_v1_literature.py | 105 ++++++++++++++++++++++++++++ 4 files changed, 289 insertions(+), 55 deletions(-) diff --git a/app/core/config.py b/app/core/config.py index 1fa2ec4..81999cb 100644 --- a/app/core/config.py +++ b/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. diff --git a/app/routes/v1.py b/app/routes/v1.py index 52d6cf8..84ac040 100644 --- a/app/routes/v1.py +++ b/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, ) diff --git a/app/services/v1_store.py b/app/services/v1_store.py index 8c8a6e5..0882b90 100644 --- a/app/services/v1_store.py +++ b/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: diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index 06124f2..bff4451 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -329,6 +329,111 @@ async def expire(self, key, ttl): assert resp.headers["X-RateLimit-Limit"] == "100" +def test_token_bucket_refills_a_token_after_the_rate_elapses(client, redis, upstream, monkeypatch): + """Design audit gap #7: the fixed one-minute window's hard reset at + :00 let a caller spend its whole budget in the last second of one + window and again in the first second of the next -- 2x limit in + under two seconds. A token bucket instead earns back one token at a + time, continuously: here, waiting exactly 1/rate seconds after + exhausting a limit-of-2 bucket earns back exactly one token, no + more.""" + import app.services.v1_store as v1_store + + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 2) + now = [1_000_000.0] + monkeypatch.setattr(v1_store.time, "time", lambda: now[0]) + + assert _post(client).status_code == 200 + assert _post(client).status_code == 200 + denied = _post(client) + assert denied.status_code == 429 + + rate = 2 / 60.0 + now[0] += 1.0 / rate # exactly enough time for one token to refill + allowed = _post(client) + assert allowed.status_code == 200 + immediately_after = _post(client) + assert immediately_after.status_code == 429 # only one token refilled, not a full reset + + +def test_rate_limit_of_zero_reports_a_fixed_reset_without_dividing_by_zero(client, redis, upstream, monkeypatch): + """limit=0 means rate=0 tokens/sec -- the reset-time computation must + not attempt to divide by that rate.""" + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 100) + redis.kv["gateway:v1:quota:42:ratelimit"] = 0 + + resp = _post(client) + assert resp.status_code == 429 + assert int(resp.headers["Retry-After"]) >= 1 + + +# --------------------------------------------------------------------------- +# Concurrent-answer limits (design audit gap #7's other remaining bullet) +# --------------------------------------------------------------------------- + + +def test_concurrency_limit_rejects_once_the_keys_slot_cap_is_reached(client, redis, upstream, monkeypatch): + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 2) + redis.kv["gateway:v1:conc:api_key:7"] = 2 # already at the cap + + resp = _post(client, idem="conc-1") + assert resp.status_code == 429 + assert resp.json()["error"]["type"] == "concurrency_limit_exceeded" + upstream.assert_not_called() + assert _usage(redis) == [] + # The denied attempt must not have left the counter net-incremented. + assert redis.kv["gateway:v1:conc:api_key:7"] == 2 + + +def test_concurrency_limit_is_enforced_per_organization_too(client, redis, upstream, monkeypatch): + """A key with no in-flight calls of its own is still blocked once its + organization's shared concurrency budget is exhausted by other + keys -- the same per-key-and-per-org pairing the rate limiter uses.""" + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 2) + redis.kv["gateway:v1:conc:org:42"] = 2 # org's shared budget already full + # This key's own counter is fresh (zero/absent). + + resp = _post(client, idem="conc-2") + assert resp.status_code == 429 + assert resp.json()["error"]["type"] == "concurrency_limit_exceeded" + # The key's own slot, acquired before the org check failed, must have + # been released back out rather than left incremented. + assert redis.kv.get("gateway:v1:conc:api_key:7", 0) == 0 + + +def test_concurrency_slot_is_released_after_a_successful_request(client, redis, upstream, monkeypatch): + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 5) + resp = _post(client) + assert resp.status_code == 200 + assert redis.kv["gateway:v1:conc:api_key:7"] == 0 + assert redis.kv["gateway:v1:conc:org:42"] == 0 + + +def test_concurrency_slot_is_released_after_a_quota_denial(client, redis, upstream, monkeypatch): + """The concurrency slot is acquired before the quota check -- a 402 + must still release it, or a key that's merely out of quota would + also look permanently "busy" to the concurrency limiter.""" + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 5) + redis.kv["gateway:v1:quota:42:literature.answer"] = 0 + + resp = _post(client, idem="conc-3") + assert resp.status_code == 402 + assert redis.kv["gateway:v1:conc:api_key:7"] == 0 + assert redis.kv["gateway:v1:conc:org:42"] == 0 + + +def test_literature_search_has_no_concurrency_cap(client, redis, upstream, monkeypatch): + """Design audit gap #7 names /v1/literature/answers specifically (the + expensive, LLM-invoking call) -- /search is retrieval-only and + carries no concurrency limit of its own.""" + monkeypatch.setattr(Config, "V1_MAX_CONCURRENT_ANSWERS", 1) + redis.kv["gateway:v1:conc:api_key:7"] = 1_000_000 # would deny /answers outright + + headers = {"Authorization": f"Bearer {KEY}"} + resp = client.post("/v1/literature/search", json={"question": "q"}, headers=headers) + assert resp.status_code == 200 + + def test_quota_exhausted_returns_402_without_calling_upstream(client, redis, upstream): redis.kv["gateway:v1:quota:42:literature.answer"] = 0 resp = _post(client, idem="q-1")