diff --git a/src/serving/api/auth/manager.py b/src/serving/api/auth/manager.py index dbc68e3..0bc8ea6 100644 --- a/src/serving/api/auth/manager.py +++ b/src/serving/api/auth/manager.py @@ -421,6 +421,24 @@ async def check_rate_limit(self, tenant_key: TenantKey) -> tuple[bool, int, int] return True, max(0, tenant_key.rate_limit_rpm - len(window)), reset_at return is_allowed, remaining, reset_at + async def charge_rate_limit(self, tenant_key: TenantKey, units: int) -> bool: + """Debit ``units`` *additional* rate-limit tokens for one request that does + more than one unit of metered work — a ``/v1/batch`` of N items runs N + engine ops but the middleware only charged the single HTTP request. + + Each unit maps to one ``check_rate_limit`` call against the same bucket, + so a batch cannot bypass the per-minute budget by bundling many (notably + the expensive NL→SQL) engine ops under one token. Every unit is debited + (the request costs its full N even when over budget); returns ``False`` if + the bucket could not absorb all of them, so the caller can reject before + doing the work. (audit S-4) + """ + allowed = True + for _ in range(max(0, units)): + ok, _remaining, _reset_at = await self.check_rate_limit(tenant_key) + allowed = allowed and ok + return allowed + def is_failed_auth_limited(self, client_ip: str) -> bool: now = self.time_source() cutoff = now - FAILED_AUTH_WINDOW_SECONDS diff --git a/src/serving/api/routers/batch.py b/src/serving/api/routers/batch.py index 0e43e69..4979c64 100644 --- a/src/serving/api/routers/batch.py +++ b/src/serving/api/routers/batch.py @@ -4,7 +4,7 @@ from typing import Any, Literal import structlog -from fastapi import APIRouter, Request +from fastapi import APIRouter, HTTPException, Request from pydantic import BaseModel, Field from src.serving.api.routers.agent_query import ( @@ -186,6 +186,26 @@ async def _execute_query_item(item: BatchItem, req: Request) -> dict[str, Any]: @router.post("/batch", response_model=BatchResponse) async def batch_query(request: BatchRequest, req: Request) -> BatchResponse: + # A batch runs one engine op per item but the auth middleware only metered the + # single HTTP request, so a tenant could drive up to 20x its per-minute budget + # (concentrated on the expensive NL path) for one token. Charge the remaining + # items against the same rate-limit bucket and reject the whole batch if the + # budget cannot absorb them. Skipped when auth is disabled (no tenant_key). + # (audit S-4) + tenant_key = getattr(req.state, "tenant_key", None) + auth_manager = getattr(req.app.state, "auth_manager", None) + extra_units = len(request.requests) - 1 + if tenant_key is not None and auth_manager is not None and extra_units > 0: + within_budget = await auth_manager.charge_rate_limit(tenant_key, extra_units) + if not within_budget: + raise HTTPException( + status_code=429, + detail=( + f"Rate limit exceeded: a {len(request.requests)}-item batch costs " + f"{len(request.requests)} of {tenant_key.rate_limit_rpm} requests/minute." + ), + ) + started_at = time.monotonic() outcomes = await asyncio.gather( *[_execute_item(item, req) for item in request.requests], diff --git a/tests/integration/test_batch.py b/tests/integration/test_batch.py index 958be29..b9b02e9 100644 --- a/tests/integration/test_batch.py +++ b/tests/integration/test_batch.py @@ -167,6 +167,40 @@ def test_batch_rejects_more_than_twenty_requests(client): assert response.status_code == 422 +def test_batch_meters_each_item_against_rate_limit(client): + # (audit S-4) a batch of N items runs N engine ops, so it must cost N + # rate-limit tokens, not one. A batch that exceeds the tenant's per-minute + # budget is rejected (429) instead of bypassing it ~20x on the expensive path. + manager = client.app.state.auth_manager + key = "batch-rpm-key" + manager.keys_by_value = { + key: TenantKey( + key=key, + name="batch-agent", + tenant="default", + rate_limit_rpm=3, + allowed_entity_types=None, + created_at=datetime.now(UTC).date(), + ) + } + manager._rate_windows.clear() + manager.rate_limiter._windows.clear() + + response = client.post( + "/v1/batch", + headers={"X-API-Key": key}, + json={ + "requests": [ + {"id": f"m{index}", "type": "metric", "params": {"name": "x", "window": "1h"}} + for index in range(5) # 1 (request) + 4 extra = 5 > rpm 3 + ] + }, + ) + + assert response.status_code == 429 + assert "Rate limit exceeded" in response.json()["detail"] + + def test_batch_requires_api_key_when_auth_is_configured(client): _set_auth(client) diff --git a/tests/unit/test_auth_manager_mutation.py b/tests/unit/test_auth_manager_mutation.py index c559bc7..97ca1bd 100644 --- a/tests/unit/test_auth_manager_mutation.py +++ b/tests/unit/test_auth_manager_mutation.py @@ -561,6 +561,26 @@ def test_bucket_key_is_deterministic_and_per_key_distinct(self) -> None: assert a1 == a2 assert a1 != b + @pytest.mark.asyncio + async def test_charge_rate_limit_within_budget_returns_true(self) -> None: + # A batch of 4 items (3 extra units) fits under a 5/min budget. (audit S-4) + m = _build_manager() + assert await m.charge_rate_limit(_key(rate_limit_rpm=5), 3) is True + + @pytest.mark.asyncio + async def test_charge_rate_limit_over_budget_returns_false(self) -> None: + # 5 extra units cannot fit a 2/min budget -> the batch is rejected. + m = _build_manager() + assert await m.charge_rate_limit(_key(rate_limit_rpm=2), 5) is False + + @pytest.mark.asyncio + async def test_charge_rate_limit_zero_units_is_noop(self) -> None: + # A single-item batch (0 extra units) debits nothing beyond the request. + m = _build_manager() + tenant_key = _key(rate_limit_rpm=1) + assert await m.charge_rate_limit(tenant_key, 0) is True + assert await m.charge_rate_limit(tenant_key, 1) is True # budget untouched + @pytest.mark.asyncio async def test_check_rate_limit_applies_local_window_when_redis_reports_full(self) -> None: m = _build_manager(