From 2267a69da5a94643c26fa324748db6ba896991e5 Mon Sep 17 00:00:00 2001
From: peterschmidt85 <andrey.cheptsov@gmail.com>
Date: Tue, 14 Oct 2025 20:16:35 -0700
Subject: [PATCH] [Internal] Extend `_CREATE_USER_HOOKS` with an optional
 config

---
 src/dstack/_internal/core/models/users.py     | 8 ++++++++
 src/dstack/_internal/server/services/users.py | 8 ++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/dstack/_internal/core/models/users.py b/src/dstack/_internal/core/models/users.py
index d09d032bf..99fb8823e 100644
--- a/src/dstack/_internal/core/models/users.py
+++ b/src/dstack/_internal/core/models/users.py
@@ -40,3 +40,11 @@ class UserTokenCreds(CoreModel):
 class UserWithCreds(User):
     creds: UserTokenCreds
     ssh_private_key: Optional[str] = None
+
+
+class UserHookConfig(CoreModel):
+    """
+    This class can be inherited to extend the user creation configuration passed to the hooks.
+    """
+
+    pass
diff --git a/src/dstack/_internal/server/services/users.py b/src/dstack/_internal/server/services/users.py
index f4c57425b..12fb0e00c 100644
--- a/src/dstack/_internal/server/services/users.py
+++ b/src/dstack/_internal/server/services/users.py
@@ -12,6 +12,7 @@
 from dstack._internal.core.models.users import (
     GlobalRole,
     User,
+    UserHookConfig,
     UserPermissions,
     UserTokenCreds,
     UserWithCreds,
@@ -79,6 +80,7 @@ async def create_user(
     email: Optional[str] = None,
     active: bool = True,
     token: Optional[str] = None,
+    config: Optional[UserHookConfig] = None,
 ) -> UserModel:
     validate_username(username)
     user_model = await get_user_model_by_name(session=session, username=username, ignore_case=True)
@@ -101,7 +103,7 @@ async def create_user(
     session.add(user)
     await session.commit()
     for func in _CREATE_USER_HOOKS:
-        await func(session, user)
+        await func(session, user, config)
     return user
 
 
@@ -267,7 +269,9 @@ def is_valid_username(username: str) -> bool:
 _CREATE_USER_HOOKS = []
 
 
-def register_create_user_hook(func: Callable[[AsyncSession, UserModel], Awaitable[None]]):
+def register_create_user_hook(
+    func: Callable[[AsyncSession, UserModel, Optional[UserHookConfig]], Awaitable[None]],
+):
     _CREATE_USER_HOOKS.append(func)