feat(BA-7365): read the rate limit the manager publishes

The limit is no longer carried in the session token, so the middleware reads
the value the manager publishes for the user. A user with no published limit
passes through: the manager republishes on every authorized request, so the
value is missing only before a user's first proxied request or after a window
of inactivity, and the manager-side limiter covers both.

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

Author: jopemachine

src/ai/backend/web/ratelimit.py

@@ -1,11 +1,15 @@
 """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.
+(``manager/api/rest/ratelimit``): the counter is keyed by the login user, so
+holding multiple keypairs does not multiply the allowance, and the limit is
+the value the manager publishes for that user. Over-limit requests are
+rejected with HTTP 429 at the web server, before they reach the manager.
+
+A user with no published limit is not limited here. The manager republishes
+on every authorized request, so the value is missing only before a user's
+first proxied request or after a window of inactivity, and the manager-side
+limiter covers both.
 """
 
 from __future__ import annotations
@@ -42,22 +46,25 @@ async def rate_limit_middleware(
     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.
+        # A session stored before the login handler started keeping the user id.
         return await handler(request)
-    rate_limit = token.get("rate_limit")
+    user_id = UserID(uuid.UUID(raw_user_id))
 
     valkey_client: ValkeyRateLimitClient = request.app["valkey_rate_limit"]
+    rate_limit = await valkey_client.get_user_rate_limit(user_id)
+    if rate_limit is None:
+        return await handler(request)
+
     rolling_count = await valkey_client.execute_rate_limit_logic(
-        user_id=UserID(uuid.UUID(raw_user_id)),
+        user_id=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-Remaining": str(max(rate_limit - rolling_count, 0)),
         "X-RateLimit-Window": str(_rlim_window),
     }
-    if rate_limit is not None and rolling_count > rate_limit:
+    if rolling_count > rate_limit:
         return web.HTTPTooManyRequests(
             text=json.dumps({
                 "type": "https://api.backend.ai/probs/rate-limit-exceeded",

tests/unit/webserver/test_ratelimit.py

@@ -17,11 +17,13 @@
 
 _USER_ID = UserID(uuid.UUID("12345678-1234-5678-1234-567812345678"))
 _RATE_LIMIT = 1000
+_AUTHENTICATED_SESSION = {"authenticated": True, "token": {"user_id": str(_USER_ID)}}
 
 
 @pytest.fixture
 def mock_valkey_rate_limit_client() -> AsyncMock:
     client = AsyncMock(spec=ValkeyRateLimitClient)
+    client.get_user_rate_limit.return_value = _RATE_LIMIT
     client.execute_rate_limit_logic.return_value = 1
     return client
 
@@ -41,6 +43,7 @@ class _PassThroughCase:
     id: str
     path: str
     session: dict[str, Any]
+    published_rate_limit: int | None = _RATE_LIMIT
 
 
 @pytest.mark.parametrize(
@@ -49,21 +52,24 @@ class _PassThroughCase:
         _PassThroughCase(
             id="non-proxied-path",
             path="/server/login",
-            session={
-                "authenticated": True,
-                "token": {"user_id": str(_USER_ID), "rate_limit": _RATE_LIMIT},
-            },
+            session=_AUTHENTICATED_SESSION,
         ),
         _PassThroughCase(
             id="unauthenticated",
             path="/func/session",
             session={},
         ),
         _PassThroughCase(
-            id="no-user-id-in-token",
+            id="session-stored-before-user-id",
             path="/func/session",
             session={"authenticated": True, "token": {"access_key": "AKTEST"}},
         ),
+        _PassThroughCase(
+            id="no-published-rate-limit",
+            path="/func/session",
+            session=_AUTHENTICATED_SESSION,
+            published_rate_limit=None,
+        ),
     ],
     ids=lambda case: case.id,
 )
@@ -75,6 +81,7 @@ async def test_pass_through_without_rate_limiting(
     handler_response: web.Response,
 ) -> None:
     mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=case.session))
+    mock_valkey_rate_limit_client.get_user_rate_limit.return_value = case.published_rate_limit
     request = make_mocked_request(
         "GET", case.path, app={"valkey_rate_limit": mock_valkey_rate_limit_client}
     )
@@ -93,18 +100,15 @@ async def test_counts_request_per_user_and_sets_headers(
     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))
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=_AUTHENTICATED_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.get_user_rate_limit.assert_awaited_once_with(_USER_ID)
     mock_valkey_rate_limit_client.execute_rate_limit_logic.assert_awaited_once_with(
         user_id=_USER_ID,
         window=ratelimit._rlim_window,
@@ -119,11 +123,7 @@ async def test_rejects_over_limit_request_with_429(
     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))
+    mocker.patch.object(ratelimit, "get_session", AsyncMock(return_value=_AUTHENTICATED_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}
@@ -136,27 +136,3 @@ async def test_rejects_over_limit_request_with_429(
     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"