Skip to content

refactor(BA-3587): Define DTO types for auth handler's soft migration #8353

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kwonkwonn wants to merge 4 commits into
base: main
Choose a base branch
from
+223 −10
Open

refactor(BA-3587): Define DTO types for auth handler's soft migration #8353

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9b06817
refactor: Define DTO types for response migration
kwonkwonn Jan 27, 2026
00a8e4c
refactor: enhance action-service for auth api
kwonkwonn Jan 27, 2026
b5e471f
refactor: modify existing auth.py to work the same with before
kwonkwonn Jan 28, 2026
e33f3d0
fix lint errors and type error from test file
kwonkwonn Jan 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/8353.enhance.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
On perpose of soft merging for manager's auth apis, defining DTO's and actions, moving auth into seperated directory.
8 changes: 8 additions & 0 deletions src/ai/backend/common/dto/manager/auth/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
"""
Common DTOs for auth system used by both Client SDK and Manager.

Import directly from submodules:
- types: AuthTokenType, AuthResponseType, TwoFactorType, AuthResponse, etc.
- request: AuthorizeRequest, SignupRequest, SignoutRequest, etc.
- response: AuthorizeResponse, SignupResponse, SignoutResponse, etc.
"""
105 changes: 105 additions & 0 deletions src/ai/backend/common/dto/manager/auth/request.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
"""
Request DTOs for auth system.
Shared between Client SDK and Manager API.
"""

from __future__ import annotations

from typing import Optional
from uuid import UUID

from pydantic import AliasChoices, Field

from ai.backend.common.api_handlers import BaseRequestModel

from .types import AuthTokenType

__all__ = (
"AuthorizeRequest",
"GetRoleRequest",
"SignupRequest",
"SignoutRequest",
"UpdateFullNameRequest",
"UpdatePasswordRequest",
"UpdatePasswordNoAuthRequest",
"UploadSSHKeypairRequest",
"VerifyAuthRequest",
)


class AuthorizeRequest(BaseRequestModel):
"""Request to authorize a user."""

type: AuthTokenType = Field(description="Authentication type (keypair or jwt)")
domain: str = Field(description="Domain name")
username: str = Field(description="Username or email")
password: str = Field(description="Password")
stoken: Optional[str] = Field(
default=None,
description="Session token",
validation_alias=AliasChoices("stoken", "sToken"),
)


class GetRoleRequest(BaseRequestModel):
"""Request to get user role."""

group: Optional[UUID] = Field(default=None, description="Group ID to check role for")


class SignupRequest(BaseRequestModel):
"""Request to sign up a new user."""

domain: str = Field(description="Domain name")
email: str = Field(description="Email address")
password: str = Field(description="Password")
username: Optional[str] = Field(default=None, description="Username")
full_name: Optional[str] = Field(default=None, description="Full name")
description: Optional[str] = Field(default=None, description="Description")


class SignoutRequest(BaseRequestModel):
"""Request to sign out a user."""

email: str = Field(
description="Email address",
validation_alias=AliasChoices("email", "username"),
)
password: str = Field(description="Password")


class UpdateFullNameRequest(BaseRequestModel):
"""Request to update user's full name."""

email: str = Field(description="Email address")
full_name: str = Field(description="New full name")


class UpdatePasswordRequest(BaseRequestModel):
"""Request to update password (authenticated)."""

old_password: str = Field(description="Current password")
new_password: str = Field(description="New password")
new_password2: str = Field(description="New password confirmation")


class UpdatePasswordNoAuthRequest(BaseRequestModel):
"""Request to update password without authentication."""

domain: str = Field(description="Domain name")
username: str = Field(description="Username or email")
current_password: str = Field(description="Current password")
new_password: str = Field(description="New password")


class UploadSSHKeypairRequest(BaseRequestModel):
"""Request to upload SSH keypair."""

pubkey: str = Field(description="SSH public key")
privkey: str = Field(description="SSH private key")


class VerifyAuthRequest(BaseRequestModel):
"""Request to verify authentication."""

echo: str = Field(description="Echo string for auth verification")
84 changes: 84 additions & 0 deletions src/ai/backend/common/dto/manager/auth/response.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
"""
Response DTOs for auth system.
Shared between Client SDK and Manager API.
"""

from __future__ import annotations

from typing import Optional

from pydantic import Field

from ai.backend.common.api_handlers import BaseResponseModel

from .types import AuthSuccessResponse

__all__ = (
"AuthorizeResponse",
"GetRoleResponse",
"SignupResponse",
"SignoutResponse",
"UpdateFullNameResponse",
"UpdatePasswordResponse",
"UpdatePasswordNoAuthResponse",
"GetSSHKeypairResponse",
"SSHKeypairResponse",
)


class AuthorizeResponse(BaseResponseModel):
"""Response for authorization."""

data: AuthSuccessResponse = Field(description="Authorization result data")


class GetRoleResponse(BaseResponseModel):
"""Response for get role."""

global_role: str = Field(description="Global role")
domain_role: str = Field(description="Domain role")
group_role: Optional[str] = Field(default=None, description="Group role")


class SignupResponse(BaseResponseModel):
"""Response for signup."""

access_key: str = Field(description="Generated access key")
secret_key: str = Field(description="Generated secret key")


class SignoutResponse(BaseResponseModel):
"""Response for signout (empty response)."""

pass


class UpdateFullNameResponse(BaseResponseModel):
"""Response for update full name (empty response)."""

pass


class UpdatePasswordResponse(BaseResponseModel):
"""Response for update password."""

error_msg: Optional[str] = Field(default=None, description="Error message if failed")


class UpdatePasswordNoAuthResponse(BaseResponseModel):
"""Response for update password without auth."""

password_changed_at: str = Field(description="Timestamp when password was changed (ISO 8601)")


class GetSSHKeypairResponse(BaseResponseModel):
"""Response for get SSH keypair (public key only)."""

ssh_public_key: str = Field(description="SSH public key")


class SSHKeypairResponse(BaseResponseModel):
"""Response for generate/upload SSH keypair (both keys)."""

ssh_public_key: str = Field(description="SSH public key")
ssh_private_key: str = Field(description="SSH private key")
27 changes: 21 additions & 6 deletions .../backend/common/dto/manager/auth/field.py → .../backend/common/dto/manager/auth/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,23 +1,38 @@
"""
Common types for auth system.
"""

from __future__ import annotations

import enum
from enum import StrEnum
from typing import Any, Self

from pydantic import BaseModel

__all__ = (
"AuthTokenType",
"AuthResponseType",
"TwoFactorType",
"AuthResponse",
"AuthSuccessResponse",
"RequireTwoFactorRegistrationResponse",
"RequireTwoFactorAuthResponse",
"parse_auth_response",
)


class AuthTokenType(enum.StrEnum):
class AuthTokenType(StrEnum):
KEYPAIR = "keypair"
JWT = "jwt"


class AuthResponseType(enum.StrEnum):
class AuthResponseType(StrEnum):
SUCCESS = "success"
REQUIRE_TWO_FACTOR_REGISTRATION = "REQUIRE_TWO_FACTOR_REGISTRATION"
REQUIRE_TWO_FACTOR_AUTH = "REQUIRE_TWO_FACTOR_AUTH"


class TwoFactorType(enum.StrEnum):
class TwoFactorType(StrEnum):
TOTP = "TOTP"


Expand Down Expand Up @@ -57,12 +72,12 @@ def to_dict(self) -> dict[str, str]:

def parse_auth_response(data: dict[str, Any]) -> AuthResponse:
raw_response_type = data.get("response_type")
respones_type = (
response_type = (
AuthResponseType(raw_response_type)
if raw_response_type is not None
else AuthResponseType.SUCCESS
)
match respones_type:
match response_type:
case AuthResponseType.SUCCESS:
return AuthSuccessResponse.model_validate(data)
case AuthResponseType.REQUIRE_TWO_FACTOR_REGISTRATION:
Expand Down
2 changes: 1 addition & 1 deletion src/ai/backend/manager/api/auth.py
View file
Open in desktop
Copy link
Collaborator

@HyeockJinKim HyeockJinKim Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Have you checked if moving the directory might cause issues? Also, it doesn't seem necessary to change the existing handlers this time. I'd like to clarify the current work direction, and I think declaring it deprecated is too hasty a decision. To mark it deprecated, we need to have a new alternative API already in place to guide users.

Copy link
Contributor Author

@kwonkwonn kwonkwonn Jan 28, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

def _init_subapp(
    pkg_name: str,
    root_app: web.Application,
    subapp: web.Application,
    global_middlewares: Iterable[Middleware],
) -> None:
    subapp.on_response_prepare.append(on_prepare)

    async def _set_root_ctx(subapp: web.Application):
        # Allow subapp's access to the root app properties.
        # These are the public APIs exposed to plugins as well.
        subapp["_root.context"] = root_app["_root.context"]
        subapp["_root_app"] = root_app

    # We must copy the public interface prior to all user-defined startup signal handlers.
    subapp.on_startup.insert(0, _set_root_ctx)
    if "prefix" not in subapp:
        subapp["prefix"] = pkg_name.split(".")[-1].replace("_", "-")
    prefix = subapp["prefix"]
    root_app.add_subapp("/" + prefix, subapp)
    root_app.middlewares.extend(global_middlewares)

when i had a investigate before migrating,
I thought it could be alright if we retain sub-servers name the same.

But i understand the concern, and it seems subapp would detect ./auth.py faster than /auth/init.py
it will be much more safer to migrate into legacy with handler/adapter pattern fully assured.

thanks you!

Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from ai.backend.common import validators as tx
from ai.backend.common.contexts.user import with_user
from ai.backend.common.data.user.types import UserData, UserRole
from ai.backend.common.dto.manager.auth.field import (
from ai.backend.common.dto.manager.auth.types import (
AuthResponseType,
AuthSuccessResponse,
AuthTokenType,
Expand Down
2 changes: 1 addition & 1 deletion src/ai/backend/manager/services/auth/actions/authorize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

from aiohttp import web

from ai.backend.common.dto.manager.auth.field import AuthTokenType
from ai.backend.common.dto.manager.auth.types import AuthTokenType
from ai.backend.manager.actions.action import BaseActionResult
from ai.backend.manager.data.auth.types import AuthorizationResult
from ai.backend.manager.services.auth.actions.base import AuthAction
Expand Down
2 changes: 1 addition & 1 deletion src/ai/backend/manager/services/auth/service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from aiohttp import web
from sqlalchemy import RowMapping

from ai.backend.common.dto.manager.auth.field import AuthTokenType
from ai.backend.common.dto.manager.auth.types import AuthTokenType
from ai.backend.common.exception import InvalidAPIParameters
from ai.backend.common.plugin.hook import ALL_COMPLETED, FIRST_COMPLETED, PASSED, HookPluginContext
from ai.backend.logging.utils import BraceStyleAdapter
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/manager/services/auth/test_authorize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import pytest
from aiohttp import web

from ai.backend.common.dto.manager.auth.field import AuthTokenType
from ai.backend.common.dto.manager.auth.types import AuthTokenType
from ai.backend.common.exception import InvalidAPIParameters
from ai.backend.common.plugin.hook import HookPluginContext, HookResult, HookResults
from ai.backend.manager.config.provider import ManagerConfigProvider
Expand Down
Loading