feat(BA-7365): key the rate limit counter by user id

`execute_rate_limit_logic()` and `get_rolling_count()` take a `UserID` and
build the counter key themselves, and the manager's rate limit middleware
counts against the authenticated user instead of the access key, so holding
multiple keypairs no longer multiplies a user's allowance.

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

Author: jopemachine

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,
@@ -44,19 +45,19 @@
 
 
 _RATE_LIMIT_SCRIPT: Final[str] = """
-local access_key = KEYS[1]
+local key = KEYS[1]
 local now = tonumber(ARGV[1])
 local window = tonumber(ARGV[2])
 local request_id = tonumber(redis.call('INCR', '__request_id'))
 if request_id >= 1e12 then
     redis.call('SET', '__request_id', 1)
 end
-if redis.call('EXISTS', access_key) == 1 then
-    redis.call('ZREMRANGEBYSCORE', access_key, 0, now - window)
+if redis.call('EXISTS', key) == 1 then
+    redis.call('ZREMRANGEBYSCORE', key, 0, now - window)
 end
-redis.call('ZADD', access_key, now, tostring(request_id))
-redis.call('EXPIRE', access_key, window)
-return redis.call('ZCARD', access_key)
+redis.call('ZADD', key, now, tostring(request_id))
+redis.call('EXPIRE', key, window)
+return redis.call('ZCARD', key)
 """
 
 
@@ -110,14 +111,14 @@ async def close(self) -> None:
     @valkey_rate_limit_resilience.apply()
     async def execute_rate_limit_logic(
         self,
-        access_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 access_key: The access key to rate limit.
+        :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.
         """
@@ -127,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=[access_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, access_key: str) -> int:
+    async def get_rolling_count(self, user_id: UserID) -> int:
         """
-        Get the current rolling count for an access key.
+        Get the current rolling count of a user.
 
-        :param access_key: The access 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(access_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

@@ -38,9 +38,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(
-                access_key=access_key,
+                user_id=request["user"]["uuid"],
                 window=_rlim_window,
             )
             if rate_limit is not None and rolling_count > rate_limit:

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(
-        access_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(
-        access_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(
-            access_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(
-            access_key="AKIAIOSFODNN7EXAMPLE",
+            user_id=_USER_ID,
             window=_rlim_window,
         )