Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/serving/api/auth/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
22 changes: 21 additions & 1 deletion src/serving/api/routers/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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],
Expand Down
34 changes: 34 additions & 0 deletions tests/integration/test_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
20 changes: 20 additions & 0 deletions tests/unit/test_auth_manager_mutation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down