From a0370de8f10459de7427e74efbae321b43b8f8c4 Mon Sep 17 00:00:00 2001 From: man4ish Date: Sat, 3 Oct 2026 18:58:44 -0500 Subject: [PATCH] feat: enforce /v1 rate limits per organization, not just per key Part of M8 (design audit gap #7's third bullet: "it is enforced per caller/key, not both per key and per organisation"). An organization holding several API keys previously got an independent rate-limit budget per key -- N keys meant N times the effective limit, since each key's counter was tracked in total isolation. _rate_limited now hits two counters per request: the existing per-subject (API key or user session) one, and a new one keyed by organization_id. Both share the same limit value (the plan-aware override from the companion M7 work, or the global default); the request is rejected if either is exhausted. The response headers report whichever counter is actually binding (the smaller remaining count), so a key blocked purely by its organization's exhausted shared budget sees 0 remaining, not its own still-fresh per-key count. Tests: 363 passed. ruff clean. New coverage: two different keys in the same organization share one budget (a key with zero requests of its own is still blocked once a sibling key exhausts the shared budget), the organization counter doesn't leak across organizations, and the reported headers reflect the binding counter correctly. Co-Authored-By: Claude Sonnet 5 --- app/routes/v1.py | 32 ++++++++++++++++++++++--- tests/test_v1_literature.py | 48 +++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/app/routes/v1.py b/app/routes/v1.py index 557106b..52d6cf8 100644 --- a/app/routes/v1.py +++ b/app/routes/v1.py @@ -5,7 +5,9 @@ against. On top of the middleware chain (auth incl. omni_sk_ API keys, policy, audit) every /v1 call gets: -- a per-caller rate limit (X-RateLimit-* headers, 429 + Retry-After), +- 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, - an org-level quota check maintained by omnibioai-billing (402), - optional Idempotency-Key replay, so a retried request is never run or billed twice, @@ -66,10 +68,34 @@ async def _rate_limited(request: Request, subject: str, org_id: str, request_id: (published by omnibioai-billing's gateway_quota_sync_service.py) when one is set, falling back to the configured global default -- `is not None`, not `or`, since a plan-specific limit of exactly 0 - is a real (if unusual) value, not "unset".""" + is a real (if unusual) value, not "unset". + + Enforced both per caller (subject -- an API key or a user session) + and per organization (design audit gap #7: "it is enforced per + caller/key, not both per key and per organisation") -- without the + organization-wide counter, an organization holding N API keys could + spread requests across them to multiply its effective limit by N, + since each key previously got its own independent budget. Both + counters share the same limit value; the request is rejected if + either is exhausted, and the headers report whichever one is + actually binding (the smaller remaining count) so the caller can + tell which budget they're hitting. + """ org_limit = await store.rate_limit_for_org(org_id) if org_id else None limit = org_limit if org_limit is not None else Config.V1_RATE_LIMIT_PER_MINUTE - allowed, remaining, reset = await store.hit_rate_limit(subject, limit) + + key_allowed, key_remaining, key_reset = await store.hit_rate_limit(subject, limit) + if org_id: + org_allowed, org_remaining, org_reset = await store.hit_rate_limit(f"org:{org_id}", limit) + else: + # No organization on this identity at all -- shouldn't normally + # happen for an authenticated /v1 caller, but fails open here + # (only the per-key check applies) rather than crashing on it. + org_allowed, org_remaining, org_reset = True, key_remaining, key_reset + + allowed = key_allowed and org_allowed + remaining, reset = (org_remaining, org_reset) if org_remaining <= key_remaining else (key_remaining, key_reset) + headers = { "X-RateLimit-Limit": str(limit), "X-RateLimit-Remaining": str(remaining), diff --git a/tests/test_v1_literature.py b/tests/test_v1_literature.py index 37aeb31..06124f2 100644 --- a/tests/test_v1_literature.py +++ b/tests/test_v1_literature.py @@ -222,6 +222,54 @@ def test_rate_limit(client, redis, upstream, monkeypatch): assert len(_usage(redis)) == 2 +def test_rate_limit_is_shared_across_multiple_keys_in_the_same_organization(client, redis, upstream, monkeypatch): + """Design audit gap #7: rate limiting must be enforced both per key + and per organization. Two different API keys belonging to the same + organization must share one organization-wide budget -- a key that + has made zero requests of its own must still be blocked once the + organization's shared budget is exhausted by a *different* key.""" + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 2) + user_key_b = {**USER, "api_key_id": 8} + + assert _post(client).status_code == 200 + assert _post(client).status_code == 200 # key A alone has now used the org's shared budget of 2 + + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=user_key_b)): + third = _post(client) # key B, same org, zero requests of its own + assert third.status_code == 429 + assert third.json()["error"]["type"] == "rate_limit_exceeded" + + +def test_rate_limit_per_key_budget_is_independent_across_organizations(client, redis, upstream, monkeypatch): + """The organization-wide counter must not leak across organizations + -- exhausting org 42's shared budget must not affect a key + belonging to a different organization.""" + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 1) + other_org_user = {**USER, "org_id": "99", "api_key_id": 9} + + assert _post(client).status_code == 200 # org 42 exhausts its budget of 1 + assert _post(client).status_code == 429 + + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=other_org_user)): + resp = _post(client) # a different organization entirely + assert resp.status_code == 200 + + +def test_rate_limit_headers_report_whichever_counter_is_binding(client, redis, upstream, monkeypatch): + """A key that has made no requests of its own, blocked purely by its + organization's exhausted shared budget, must see 0 remaining -- not + its own, still-fresh per-key count.""" + monkeypatch.setattr(Config, "V1_RATE_LIMIT_PER_MINUTE", 1) + user_key_b = {**USER, "api_key_id": 8} + + assert _post(client).status_code == 200 # key A exhausts the org's shared budget of 1 + + with patch.object(_main_mod.iam, "validate_api_key", AsyncMock(return_value=user_key_b)): + resp = _post(client) + assert resp.status_code == 429 + assert resp.headers["X-RateLimit-Remaining"] == "0" + + def test_rate_limit_uses_the_organizations_plan_specific_override(client, redis, upstream, monkeypatch): """omnibioai-billing publishes this org's plan-specific override under gateway:v1:quota:{org}:ratelimit (see