feat(BA-7365): key the manager rate limiter by user id

`execute_rate_limit_logic()` and `get_rolling_count()` now take a `UserID`
and build the counter key themselves, so the manager and the web server
share one per-user allowance instead of the manager counting per access key.

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

Author: jopemachine

changes/13771.feature.md

@@ -1 +1 @@
-Add per-user rate limiting to the webserver that rejects over-limit requests with HTTP 429 before they reach the manager, enforcing the existing keypair rate limit against a single per-user allowance.
+Add per-user rate limiting to the webserver that rejects over-limit requests with HTTP 429 before they reach the manager, and key the manager-side rate limiter by the authenticated user as well, so holding multiple keypairs no longer multiplies a user's allowance.

src/ai/backend/common/clients/valkey_client/valkey_rate_limit/client.py

@@ -10,6 +10,7 @@
     create_valkey_client,
 )
 from ai.backend.common.exception import BackendAIError
+from ai.backend.common.identifier.user import UserID
 from ai.backend.common.metrics.metric import DomainType, LayerType
 from ai.backend.common.resilience import (
     BackoffStrategy,
@@ -110,15 +111,14 @@ async def close(self) -> None:
     @valkey_rate_limit_resilience.apply()
     async def execute_rate_limit_logic(
         self,
-        key: str,
+        user_id: UserID,
         window: int = _DEFAULT_RATE_LIMIT_EXPIRATION,
     ) -> int:
         """
         Execute the rate limiting logic for rolling counter.
         This replicates the Lua script logic using individual commands.
 
-        :param key: The counter key to rate limit, e.g. an access key or a
-            namespaced user key.
+        :param user_id: The user the rolling counter is keyed by.
         :param window: The time window for rate limiting in seconds.
         :return: The current count.
         """
@@ -128,23 +128,23 @@ async def execute_rate_limit_logic(
         async with self._client.client() as conn:
             result = await conn.invoke_script(
                 Script(_RATE_LIMIT_SCRIPT),
-                keys=[key],
+                keys=[f"user.{user_id}"],
                 args=[str(now_float), str(window)],
             )
 
         # The last result is the count
         return cast(int, result)
 
     @valkey_rate_limit_resilience.apply()
-    async def get_rolling_count(self, key: str) -> int:
+    async def get_rolling_count(self, user_id: UserID) -> int:
         """
-        Get the current rolling count for a counter key.
+        Get the current rolling count of a user.
 
-        :param key: The counter key to get the count for.
+        :param user_id: The user the rolling counter is keyed by.
         :return: The current count.
         """
         async with self._client.client() as conn:
-            return await conn.zcard(key)
+            return await conn.zcard(f"user.{user_id}")
 
     @valkey_rate_limit_resilience.apply()
     async def set_rate_limit_config(

src/ai/backend/manager/api/gql_legacy/keypair.py

@@ -13,6 +13,7 @@
 
 from ai.backend.common.clients.valkey_client.valkey_rate_limit.client import ValkeyRateLimitClient
 from ai.backend.common.defs import REDIS_RATE_LIMIT_DB, RedisRole
+from ai.backend.common.identifier.user import UserID
 from ai.backend.common.types import AccessKey
 from ai.backend.manager.data.kernel.types import KernelStatus
 from ai.backend.manager.data.keypair.types import KeyPairCreator, KeyPairData
@@ -196,7 +197,7 @@ async def resolve_rolling_count(self, info: graphene.ResolveInfo) -> int:
             human_readable_name="ratelimit",
         )
         try:
-            return await valkey_client.get_rolling_count(self.access_key)
+            return await valkey_client.get_rolling_count(UserID(self.user))
         finally:
             await valkey_client.close()
 

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
@@ -38,9 +43,8 @@ async def rlim_middleware(
         """Global middleware implementing a rolling-counter rate limiter."""
         if request["is_authorized"]:
             rate_limit = request["keypair"]["rate_limit"]
-            access_key = request["keypair"]["access_key"]
             rolling_count = await valkey_client.execute_rate_limit_logic(
-                key=access_key,
+                user_id=request["user"]["uuid"],
                 window=_rlim_window,
             )
             if rate_limit is not None and rolling_count > rate_limit:

src/ai/backend/web/ratelimit.py

@@ -1,23 +1,25 @@
 """Per-user rate limiting for requests proxied to the manager.
 
-The same rolling-counter mechanism 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), but the counter
-is keyed by the login user instead of the access key so that 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.
+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
 
@@ -38,15 +40,15 @@ async def rate_limit_middleware(
     if not session.get("authenticated", False):
         return await handler(request)
     token = session.get("token") or {}
-    user_id = token.get("user_id")
-    if user_id is None:
+    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(
-        f"user.{user_id}",
+        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

tests/unit/common/clients/valkey_client/test_valkey_rate_limit_client.py

@@ -1,30 +1,33 @@
 from __future__ import annotations
 
-import random
+import uuid
 
 from ai.backend.common.clients.valkey_client.valkey_rate_limit.client import (
     ValkeyRateLimitClient,
 )
+from ai.backend.common.identifier.user import UserID
 
 
 async def test_valkey_rate_limit_logic_execution(
     test_valkey_rate_limit: ValkeyRateLimitClient,
 ) -> None:
     """Test rate limiting logic execution."""
-    access_key = f"test-logic-{random.randint(1000, 9999)}"
+    user_id = UserID(uuid.uuid4())
 
     # Execute the rate limiting logic
     result = await test_valkey_rate_limit.execute_rate_limit_logic(
-        key=access_key,
+        user_id=user_id,
         window=60,
     )
 
     assert result == 1  # First request should return 1
 
     # Execute again
     result2 = await test_valkey_rate_limit.execute_rate_limit_logic(
-        key=access_key,
+        user_id=user_id,
         window=60,
     )
 
     assert result2 == 2  # Second request should return 2
+
+    assert await test_valkey_rate_limit.get_rolling_count(user_id) == 2

tests/unit/manager/api/test_ratelimit.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import uuid
 from dataclasses import dataclass
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
@@ -8,9 +9,12 @@
 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.manager.api.rest.ratelimit.handler import _rlim_window, make_rlim_middleware
 from ai.backend.manager.errors.api import RateLimitExceeded
 
+_USER_ID = UserID(uuid.UUID("12345678-1234-5678-1234-567812345678"))
+
 
 @dataclass
 class RateLimitSuccessCase:
@@ -65,16 +69,16 @@ def mock_request_anonymous(self) -> web.Request:
     def mock_request_authorized(self) -> web.Request:
         """Mock request for authorized user."""
         request = MagicMock(spec=web.Request)
-        keypair_data = {
-            "rate_limit": 30000,
-            "access_key": "AKIAIOSFODNN7EXAMPLE",
-        }
+        keypair_data = {"rate_limit": 30000}
+        user_data = {"uuid": _USER_ID}
 
         def getitem(key: Any) -> Any:
             if key == "is_authorized":
                 return True
             if key == "keypair":
                 return keypair_data
+            if key == "user":
+                return user_data
             return None
 
         request.__getitem__ = MagicMock(side_effect=getitem)
@@ -162,7 +166,7 @@ async def test_authorized_query_within_rate_limit(
 
         # Valkey should be called for authorized requests
         mock_valkey_client.execute_rate_limit_logic.assert_called_once_with(
-            key="AKIAIOSFODNN7EXAMPLE",
+            user_id=_USER_ID,
             window=_rlim_window,
         )
 
@@ -211,6 +215,6 @@ async def test_authorized_query_exceeds_rate_limit(
 
         # Valkey should still be called
         mock_valkey_client.execute_rate_limit_logic.assert_called_once_with(
-            key="AKIAIOSFODNN7EXAMPLE",
+            user_id=_USER_ID,
             window=_rlim_window,
         )

tests/unit/webserver/test_ratelimit.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import uuid
 from dataclasses import dataclass
 from typing import Any
 from unittest.mock import AsyncMock
@@ -10,10 +11,11 @@
 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 = "12345678-1234-5678-1234-567812345678"
+_USER_ID = UserID(uuid.UUID("12345678-1234-5678-1234-567812345678"))
 _RATE_LIMIT = 1000
 
 
@@ -49,7 +51,7 @@ class _PassThroughCase:
             path="/server/login",
             session={
                 "authenticated": True,
-                "token": {"user_id": _USER_ID, "rate_limit": _RATE_LIMIT},
+                "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
             },
         ),
         _PassThroughCase(
@@ -93,7 +95,7 @@ async def test_counts_request_per_user_and_sets_headers(
 ) -> None:
     session = {
         "authenticated": True,
-        "token": {"user_id": _USER_ID, "rate_limit": _RATE_LIMIT},
+        "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
     }
     mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=session))
     request = make_mocked_request(
@@ -104,7 +106,7 @@ async def test_counts_request_per_user_and_sets_headers(
 
     assert response is handler_response
     mock_valkey_rate_limit_client.execute_rate_limit_logic.assert_awaited_once_with(
-        f"user.{_USER_ID}",
+        user_id=_USER_ID,
         window=ratelimit._rlim_window,
     )
     assert response.headers["X-RateLimit-Limit"] == str(_RATE_LIMIT)
@@ -119,7 +121,7 @@ async def test_rejects_over_limit_request_with_429(
 ) -> None:
     session = {
         "authenticated": True,
-        "token": {"user_id": _USER_ID, "rate_limit": _RATE_LIMIT},
+        "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
@@ -144,7 +146,7 @@ async def test_null_rate_limit_counts_but_never_rejects(
 ) -> None:
     session = {
         "authenticated": True,
-        "token": {"user_id": _USER_ID, "rate_limit": None},
+        "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