feat(BA-7365): add a per-user rate limit middleware to the web server

Requests proxied to the manager (`/func/*`) are counted against the login
user's rolling counter and rejected with HTTP 429 once the session's rate
limit is passed, so floods no longer reach the manager. The middleware is
registered after `setup_session()` because it reads the session storage the
session middleware installs.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: jopemachine

changes/13771.feature.md

@@ -0,0 +1 @@
+Reject over-limit requests at the web server with HTTP 429 before they are proxied to the manager, counting them against the requesting user's rate limit.

src/ai/backend/manager/api/rest/ratelimit/handler.py

@@ -3,6 +3,11 @@
 This module provides the ``rlim_middleware`` function which is installed
 as a global aiohttp middleware.  There are no route handlers — rate
 limiting is applied transparently to all authorized requests.
+
+The rolling counter is keyed by the authenticated user, the same counter the
+web server's limiter uses. This is the fallback guard for clients that reach
+the manager directly; requests proxied by the web server are already limited
+there.
 """
 
 from __future__ import annotations

src/ai/backend/web/ratelimit.py

@@ -0,0 +1,72 @@
+"""Per-user rate limiting for requests proxied to the manager.
+
+The same rolling counter and limit value as the manager-side rate limiter
+(``manager/api/rest/ratelimit``): the limit is the keypair ``rate_limit``
+delivered at login (``None`` means unlimited), and the counter is keyed by the
+login user, so holding multiple keypairs does not multiply the allowance.
+Over-limit requests are rejected with HTTP 429 at the web server, before they
+reach the manager.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from collections.abc import Awaitable, Callable
+from typing import Final
+
+from aiohttp import web
+
+from ai.backend.common.clients.valkey_client.valkey_rate_limit.client import ValkeyRateLimitClient
+from ai.backend.common.identifier.user import UserID
+from ai.backend.common.web.session import get_session
+from ai.backend.logging import BraceStyleAdapter
+
+log: Final = BraceStyleAdapter(logging.getLogger(__spec__.name))
+
+_rlim_window: Final = 60 * 15
+_RATE_LIMITED_PATH_PREFIX: Final = "/func/"
+
+
+@web.middleware
+async def rate_limit_middleware(
+    request: web.Request,
+    handler: Callable[[web.Request], Awaitable[web.StreamResponse]],
+) -> web.StreamResponse:
+    if not request.path.startswith(_RATE_LIMITED_PATH_PREFIX):
+        return await handler(request)
+    session = await get_session(request)
+    if not session.get("authenticated", False):
+        return await handler(request)
+    token = session.get("token") or {}
+    raw_user_id = token.get("user_id")
+    if raw_user_id is None:
+        # Session created against a manager that does not send the user ID yet.
+        return await handler(request)
+    rate_limit = token.get("rate_limit")
+
+    valkey_client: ValkeyRateLimitClient = request.app["valkey_rate_limit"]
+    rolling_count = await valkey_client.execute_rate_limit_logic(
+        user_id=UserID(uuid.UUID(raw_user_id)),
+        window=_rlim_window,
+    )
+    remaining = max(rate_limit - rolling_count, 0) if rate_limit is not None else rolling_count
+    rlim_headers = {
+        "X-RateLimit-Limit": str(rate_limit),
+        "X-RateLimit-Remaining": str(remaining),
+        "X-RateLimit-Window": str(_rlim_window),
+    }
+    if rate_limit is not None and rolling_count > rate_limit:
+        return web.HTTPTooManyRequests(
+            text=json.dumps({
+                "type": "https://api.backend.ai/probs/rate-limit-exceeded",
+                "title": "You have reached your API query rate limit.",
+            }),
+            content_type="application/problem+json",
+            headers=rlim_headers,
+        )
+    response = await handler(request)
+    if not response.prepared:
+        response.headers.update(rlim_headers)
+    return response

src/ai/backend/web/server.py

@@ -47,8 +47,9 @@
 from ai.backend.client.v2.registry import BackendAIClientRegistry
 from ai.backend.common import config
 from ai.backend.common.clients.http_client.client_pool import ClientPool
+from ai.backend.common.clients.valkey_client.valkey_rate_limit.client import ValkeyRateLimitClient
 from ai.backend.common.clients.valkey_client.valkey_session.client import ValkeySessionClient
-from ai.backend.common.defs import REDIS_STATISTICS_DB, RedisRole
+from ai.backend.common.defs import REDIS_RATE_LIMIT_DB, REDIS_STATISTICS_DB, RedisRole
 from ai.backend.common.dto.internal.health import (
     ConnectivityCheckResponse,
     HealthResponse,
@@ -92,6 +93,7 @@
     ManagerPoolGateHealthChecker,
 )
 from ai.backend.web.config.unified import EventLoopType, ServiceMode, WebServerUnifiedConfig
+from ai.backend.web.ratelimit import rate_limit_middleware
 from ai.backend.web.security import SecurityPolicy, csp_nonce_var, security_policy_middleware
 
 from . import __version__, user_agent
@@ -857,6 +859,13 @@ async def redis_ctx(
     # Keep app["redis"] key for compatibility
     app["redis"] = valkey_session_client
 
+    valkey_rate_limit_client = await ValkeyRateLimitClient.create(
+        valkey_target=valkey_profile_target.profile_target(RedisRole.RATE_LIMIT),
+        db_id=REDIS_RATE_LIMIT_DB,
+        human_readable_name="web.ratelimit",
+    )
+    app["valkey_rate_limit"] = valkey_rate_limit_client
+
     if pidx == 0 and config.session.flush_on_startup:
         await valkey_session_client.flush_all_sessions()
         log.info("flushed session storage.")
@@ -869,9 +878,13 @@ async def redis_ctx(
         secure=config.security.cookie_secure,
     )
     setup_session(app, redis_storage)
+    # Must come after ``setup_session()``: the session middleware it appends is what
+    # populates the request with the session storage the rate limiter reads.
+    app.middlewares.append(rate_limit_middleware)
     try:
         yield
     finally:
+        await valkey_rate_limit_client.close()
         await valkey_session_client.close()
 
 

tests/unit/webserver/test_ratelimit.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+from aiohttp import web
+from aiohttp.test_utils import make_mocked_request
+from pytest_mock import MockerFixture
+
+from ai.backend.common.clients.valkey_client.valkey_rate_limit.client import ValkeyRateLimitClient
+from ai.backend.common.identifier.user import UserID
+from ai.backend.web import ratelimit
+from ai.backend.web.ratelimit import rate_limit_middleware
+
+_USER_ID = UserID(uuid.UUID("12345678-1234-5678-1234-567812345678"))
+_RATE_LIMIT = 1000
+
+
+@pytest.fixture
+def mock_valkey_rate_limit_client() -> AsyncMock:
+    client = AsyncMock(spec=ValkeyRateLimitClient)
+    client.execute_rate_limit_logic.return_value = 1
+    return client
+
+
+@pytest.fixture
+def handler_response() -> web.Response:
+    return web.Response(text="ok")
+
+
+@pytest.fixture
+def handler(handler_response: web.Response) -> AsyncMock:
+    return AsyncMock(return_value=handler_response)
+
+
+@dataclass(frozen=True)
+class _PassThroughCase:
+    id: str
+    path: str
+    session: dict[str, Any]
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        _PassThroughCase(
+            id="non-proxied-path",
+            path="/server/login",
+            session={
+                "authenticated": True,
+                "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
+            },
+        ),
+        _PassThroughCase(
+            id="unauthenticated",
+            path="/func/session",
+            session={},
+        ),
+        _PassThroughCase(
+            id="no-user-id-in-token",
+            path="/func/session",
+            session={"authenticated": True, "token": {"access_key": "AKTEST"}},
+        ),
+    ],
+    ids=lambda case: case.id,
+)
+async def test_pass_through_without_rate_limiting(
+    case: _PassThroughCase,
+    mocker: MockerFixture,
+    mock_valkey_rate_limit_client: AsyncMock,
+    handler: AsyncMock,
+    handler_response: web.Response,
+) -> None:
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=case.session))
+    request = make_mocked_request(
+        "GET", case.path, app={"valkey_rate_limit": mock_valkey_rate_limit_client}
+    )
+
+    response = await rate_limit_middleware(request, handler)
+
+    assert response is handler_response
+    handler.assert_awaited_once_with(request)
+    mock_valkey_rate_limit_client.execute_rate_limit_logic.assert_not_called()
+    assert "X-RateLimit-Limit" not in response.headers
+
+
+async def test_counts_request_per_user_and_sets_headers(
+    mocker: MockerFixture,
+    mock_valkey_rate_limit_client: AsyncMock,
+    handler: AsyncMock,
+    handler_response: web.Response,
+) -> None:
+    session = {
+        "authenticated": True,
+        "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
+    }
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=session))
+    request = make_mocked_request(
+        "GET", "/func/session", app={"valkey_rate_limit": mock_valkey_rate_limit_client}
+    )
+
+    response = await rate_limit_middleware(request, handler)
+
+    assert response is handler_response
+    mock_valkey_rate_limit_client.execute_rate_limit_logic.assert_awaited_once_with(
+        user_id=_USER_ID,
+        window=ratelimit._rlim_window,
+    )
+    assert response.headers["X-RateLimit-Limit"] == str(_RATE_LIMIT)
+    assert response.headers["X-RateLimit-Remaining"] == str(_RATE_LIMIT - 1)
+    assert response.headers["X-RateLimit-Window"] == str(ratelimit._rlim_window)
+
+
+async def test_rejects_over_limit_request_with_429(
+    mocker: MockerFixture,
+    mock_valkey_rate_limit_client: AsyncMock,
+    handler: AsyncMock,
+) -> None:
+    session = {
+        "authenticated": True,
+        "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
+    }
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=session))
+    mock_valkey_rate_limit_client.execute_rate_limit_logic.return_value = _RATE_LIMIT + 1
+    request = make_mocked_request(
+        "GET", "/func/session", app={"valkey_rate_limit": mock_valkey_rate_limit_client}
+    )
+
+    response = await rate_limit_middleware(request, handler)
+
+    handler.assert_not_called()
+    assert isinstance(response, web.HTTPTooManyRequests)
+    assert response.content_type == "application/problem+json"
+    assert response.headers["X-RateLimit-Limit"] == str(_RATE_LIMIT)
+    assert response.headers["X-RateLimit-Remaining"] == "0"
+
+
+async def test_null_rate_limit_counts_but_never_rejects(
+    mocker: MockerFixture,
+    mock_valkey_rate_limit_client: AsyncMock,
+    handler: AsyncMock,
+    handler_response: web.Response,
+) -> None:
+    session = {
+        "authenticated": True,
+        "token": {"user_id": str(_USER_ID), "rate_limit": None},
+    }
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=session))
+    mock_valkey_rate_limit_client.execute_rate_limit_logic.return_value = 10_000_000
+    request = make_mocked_request(
+        "GET", "/func/session", app={"valkey_rate_limit": mock_valkey_rate_limit_client}
+    )
+
+    response = await rate_limit_middleware(request, handler)
+
+    assert response is handler_response
+    handler.assert_awaited_once_with(request)
+    mock_valkey_rate_limit_client.execute_rate_limit_logic.assert_awaited_once()
+    assert response.headers["X-RateLimit-Remaining"] == "10000000"