feat(BA-7364): deliver the login user id in the auth response

changes/13777.feature.md

@@ -0,0 +1 @@
+Return the login user's id and rate limit from the authorization API, and keep both in the web server session.

src/ai/backend/common/dto/manager/auth/types.py

@@ -7,7 +7,9 @@
 from enum import StrEnum
 from typing import Any, Self
 
-from ai.backend.common.types import BackendAISchema
+from ai.backend.common.data.user.types import UserRole
+from ai.backend.common.identifier.user import UserID
+from ai.backend.common.types import AccessKey, BackendAISchema, SecretKey
 
 __all__ = (
     "AuthTokenType",
@@ -50,14 +52,16 @@ def parse(cls, data: dict[str, Any]) -> Self:
 class AuthSuccessResponse(AuthResponse):
     """Returned when authorization succeeds without requiring 2FA."""
 
-    access_key: str
-    secret_key: str
-    role: str
+    access_key: AccessKey
+    secret_key: SecretKey
+    role: UserRole
     status: str
     session_token: str
+    user_id: UserID
+    rate_limit: int | None
     type: AuthTokenType = AuthTokenType.KEYPAIR
 
-    def to_dict(self) -> dict[str, str]:
+    def to_dict(self) -> dict[str, Any]:
         return self.model_dump(mode="json")
 
 

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

@@ -168,6 +168,8 @@ async def authorize(
             role=auth_result.role,
             status=auth_result.status,
             session_token=auth_result.session_token,
+            user_id=auth_result.user_id,
+            rate_limit=auth_result.rate_limit,
         )
         resp = AuthorizeResponse(data=data)
         return APIResponse.build(HTTPStatus.OK, resp)

src/ai/backend/manager/data/auth/types.py

@@ -23,12 +23,13 @@ class SSHKeypair:
 
 @dataclass
 class AuthorizationResult:
-    user_id: uuid.UUID
-    access_key: str
-    secret_key: str
-    role: str
+    user_id: UserID
+    access_key: AccessKey
+    secret_key: SecretKey
+    role: UserRole
     status: str
     session_token: str
+    rate_limit: int | None
 
 
 @dataclass

src/ai/backend/manager/services/auth/service.py

@@ -18,8 +18,9 @@
 )
 from ai.backend.common.dto.manager.auth.types import AuthTokenType
 from ai.backend.common.exception import InvalidAPIParameters, UserResourcePolicyNotFound
+from ai.backend.common.identifier.user import UserID
 from ai.backend.common.plugin.hook import ALL_COMPLETED, FIRST_COMPLETED, PASSED, HookPluginContext
-from ai.backend.common.types import AccessKey, SSHPrivateKey, SSHPublicKey
+from ai.backend.common.types import AccessKey, SecretKey, SSHPrivateKey, SSHPublicKey
 from ai.backend.logging.utils import BraceStyleAdapter
 from ai.backend.manager.config.provider import ManagerConfigProvider
 from ai.backend.manager.config.unified import AuthConfig
@@ -394,7 +395,7 @@ async def _create_login_session(
                 token=LoginSessionTokenData(
                     type="keypair",
                     access_key=keypair_row.access_key,
-                    secret_key=keypair_row.secret_key or "",
+                    secret_key=keypair_row.secret_key,
                     role=user.role,
                     status=user.status,
                 ),
@@ -409,12 +410,14 @@ async def _create_login_session(
         return AuthorizeActionResult(
             stream_response=None,
             authorization_result=AuthorizationResult(
-                access_key=keypair_row.access_key,
-                secret_key=keypair_row.secret_key or "",
-                user_id=user.uuid,
-                role=user.role,
+                access_key=AccessKey(keypair_row.access_key),
+                secret_key=SecretKey(keypair_row.secret_key),
+                user_id=UserID(user.uuid),
+                role=UserRole(user.role),
                 status=user.status,
                 session_token=session_result.session_token,
+                # TODO: take this from the user resource policy instead of the keypair.
+                rate_limit=keypair_row.rate_limit,
             ),
         )
 

src/ai/backend/web/server.py

@@ -506,6 +506,8 @@ async def _set_login_history(last_login_attempt: float, login_fail_count: float
                     "secret_key": token.secret_key,
                     "role": token.role,
                     "status": token.status,
+                    "user_id": str(token.user_id),
+                    "rate_limit": token.rate_limit,
                 }
                 public_return = {
                     "access_key": token.access_key,
@@ -767,6 +769,8 @@ async def token_login_handler(request: web.Request) -> web.Response:
             "secret_key": token.secret_key,
             "role": token.role,
             "status": token.status,
+            "user_id": str(token.user_id),
+            "rate_limit": token.rate_limit,
         }
         public_return = {
             "access_key": token.access_key,

tests/unit/client_v2/test_auth_client.py

@@ -81,6 +81,8 @@ async def test_authorize(self, sample_client_type_id: UUID) -> None:
                     "role": "admin",
                     "status": "active",
                     "session_token": "test_session_token",
+                    "user_id": "12345678-1234-5678-1234-567812345678",
+                    "rate_limit": None,
                     "type": AuthTokenType.KEYPAIR,
                 },
             }

tests/unit/common/dto/manager/auth/test_auth_response.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+from uuid import uuid4
+
 from ai.backend.common.api_handlers import BaseResponseModel
+from ai.backend.common.data.user.types import UserRole
 from ai.backend.common.dto.manager.auth.response import (
     AuthorizeResponse,
     GetRoleResponse,
@@ -18,16 +21,20 @@
     AuthSuccessResponse,
     AuthTokenType,
 )
+from ai.backend.common.identifier.user import UserID
+from ai.backend.common.types import AccessKey, SecretKey
 
 
 def test_authorize_response() -> None:
     data = AuthSuccessResponse(
         response_type=AuthResponseType.SUCCESS,
-        access_key="AKIAIOSFODNN7EXAMPLE",
-        secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
-        role="user",
+        access_key=AccessKey("AKIAIOSFODNN7EXAMPLE"),
+        secret_key=SecretKey("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"),
+        role=UserRole.USER,
         status="active",
         session_token="test_session_token",
+        user_id=UserID(uuid4()),
+        rate_limit=None,
         type=AuthTokenType.KEYPAIR,
     )
     resp = AuthorizeResponse(data=data)

tests/unit/common/dto/manager/auth/test_auth_types.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from uuid import uuid4
+
+from ai.backend.common.data.user.types import UserRole
 from ai.backend.common.dto.manager.auth.types import (
     AuthResponseType,
     AuthSuccessResponse,
@@ -9,6 +12,8 @@
     TwoFactorType,
     parse_auth_response,
 )
+from ai.backend.common.identifier.user import UserID
+from ai.backend.common.types import AccessKey, SecretKey
 
 
 def test_auth_token_type_values() -> None:
@@ -34,11 +39,13 @@ def test_two_factor_type_values() -> None:
 def test_auth_success_response_creation() -> None:
     resp = AuthSuccessResponse(
         response_type=AuthResponseType.SUCCESS,
-        access_key="AKTEST",
-        secret_key="SKTEST",
-        role="user",
+        access_key=AccessKey("AKTEST"),
+        secret_key=SecretKey("SKTEST"),
+        role=UserRole.USER,
         status="active",
         session_token="test_session_token",
+        user_id=UserID(uuid4()),
+        rate_limit=None,
     )
     assert resp.access_key == "AKTEST"
     assert resp.secret_key == "SKTEST"
@@ -50,11 +57,13 @@ def test_auth_success_response_creation() -> None:
 def test_auth_success_response_to_dict() -> None:
     resp = AuthSuccessResponse(
         response_type=AuthResponseType.SUCCESS,
-        access_key="AK",
-        secret_key="SK",
-        role="admin",
+        access_key=AccessKey("AK"),
+        secret_key=SecretKey("SK"),
+        role=UserRole.ADMIN,
         status="active",
         session_token="test_session_token",
+        user_id=UserID(uuid4()),
+        rate_limit=None,
         type=AuthTokenType.JWT,
     )
     d = resp.to_dict()
@@ -115,6 +124,8 @@ def test_parse_auth_response_success() -> None:
         "role": "user",
         "status": "active",
         "session_token": "test_token",
+        "user_id": "12345678-1234-5678-1234-567812345678",
+        "rate_limit": None,
     }
     result = parse_auth_response(data)
     assert isinstance(result, AuthSuccessResponse)
@@ -150,6 +161,8 @@ def test_parse_auth_response_explicit_success() -> None:
         "role": "user",
         "status": "active",
         "session_token": "test_token",
+        "user_id": "12345678-1234-5678-1234-567812345678",
+        "rate_limit": None,
     }
     result = parse_auth_response(data)
     assert isinstance(result, AuthSuccessResponse)

tests/unit/manager/api/auth/test_handlers.py

@@ -33,6 +33,8 @@
     UploadSSHKeypairRequest,
     VerifyAuthRequest,
 )
+from ai.backend.common.identifier.user import UserID
+from ai.backend.common.types import AccessKey, SecretKey
 from ai.backend.manager.api.rest.auth.handler import AuthHandler
 from ai.backend.manager.api.rest.middleware.auth import (
     TRUSTED_PROXY_NETWORKS_KEY,
@@ -240,12 +242,13 @@ def authorize_result(self) -> AuthorizeActionResult:
         return AuthorizeActionResult(
             stream_response=None,
             authorization_result=AuthorizationResult(
-                user_id=uuid.uuid4(),
-                access_key="TESTKEY",
-                secret_key="TESTSECRET",
+                user_id=UserID(uuid.uuid4()),
+                access_key=AccessKey("TESTKEY"),
+                secret_key=SecretKey("TESTSECRET"),
                 role=UserRole.USER,
                 status=UserStatus.ACTIVE,
                 session_token="test_session_token",
+                rate_limit=30000,
             ),
         )