refactor: prepare Android ViewModel creation boundary [WPB-23495]

app/src/main/kotlin/com/wire/android/di/AssistedViewModelExt.kt

@@ -19,11 +19,36 @@ package com.wire.android.di
 
 import androidx.activity.ComponentActivity
 import androidx.activity.viewModels
+import androidx.compose.runtime.Composable
+import androidx.hilt.navigation.compose.hiltViewModel
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelStoreOwner
+import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
 import dagger.hilt.android.lifecycle.withCreationCallback
 
 inline fun <reified VM : ViewModel, reified VMF> ComponentActivity.assistedViewModels(
     crossinline create: (VMF) -> VM
 ) = viewModels<VM>(extrasProducer = {
     defaultViewModelCreationExtras.withCreationCallback<VMF> { factory -> create(factory) }
 })
+
+@Composable
+inline fun <reified VM : ViewModel> wireViewModel(
+    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
+        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
+    },
+    key: String? = null,
+): VM = hiltViewModel(viewModelStoreOwner = viewModelStoreOwner, key = key)
+
+@Composable
+inline fun <reified VM : ViewModel, reified VMF> wireViewModel(
+    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
+        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
+    },
+    key: String? = null,
+    noinline creationCallback: (VMF) -> VM
+): VM = hiltViewModel(
+    viewModelStoreOwner = viewModelStoreOwner,
+    key = key,
+    creationCallback = creationCallback
+)

app/src/main/kotlin/com/wire/android/navigation/MainNavHost.kt

@@ -31,6 +31,7 @@ import com.ramcosta.composedestinations.DestinationsNavHost
 import com.ramcosta.composedestinations.generated.app.destinations.ConversationScreenDestination
 import com.ramcosta.composedestinations.generated.app.destinations.NewLoginPasswordScreenDestination
 import com.ramcosta.composedestinations.generated.app.destinations.NewLoginVerificationCodeScreenDestination
+import com.ramcosta.composedestinations.generated.app.navArgs
 import com.ramcosta.composedestinations.generated.app.navgraphs.NewConversationGraph
 import com.ramcosta.composedestinations.generated.app.navgraphs.PersonalToTeamMigrationGraph
 import com.ramcosta.composedestinations.generated.cells.destinations.SearchScreenDestination
@@ -49,7 +50,9 @@ import com.ramcosta.composedestinations.scope.resultRecipient
 import com.ramcosta.composedestinations.spec.Direction
 import com.wire.android.feature.cells.ui.CellViewModel
 import com.wire.android.feature.sketch.model.DrawingCanvasNavBackArgs
+import com.wire.android.di.wireViewModel
 import com.wire.android.navigation.transition.LocalSharedTransitionScope
+import com.wire.android.ui.authentication.login.LoginNavArgs
 import com.wire.android.ui.authentication.login.email.LoginEmailViewModel
 import com.wire.android.ui.home.conversations.ConversationScreen
 import com.wire.android.ui.home.newconversation.NewConversationViewModel
@@ -94,7 +97,12 @@ fun MainNavHost(
                         val loginPasswordEntry = remember(navBackStackEntry) {
                             navController.getBackStackEntry(NewLoginPasswordScreenDestination.route)
                         }
-                        dependency(hiltViewModel<LoginEmailViewModel>(loginPasswordEntry))
+                        dependency(
+                            wireViewModel<LoginEmailViewModel, LoginEmailViewModel.Factory>(
+                                viewModelStoreOwner = loginPasswordEntry,
+                                creationCallback = { factory -> factory.create(loginPasswordEntry.navArgs<LoginNavArgs>()) }
+                            )
+                        )
                     }
 
                     // 👇 To reuse CellViewModel from the parent screen on SearchScreen

app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginSavedInputStore.kt

@@ -0,0 +1,24 @@
+/*
+ * Wire
+ * Copyright (C) 2026 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+
+package com.wire.android.ui.authentication.login
+
+interface LoginSavedInputStore {
+    var userIdentifier: String?
+    var ssoCode: String?
+}

app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginScreen.kt

@@ -44,12 +44,12 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalFocusManager
 import androidx.compose.ui.platform.LocalSoftwareKeyboardController
 import androidx.compose.ui.res.stringResource
-import androidx.hilt.navigation.compose.hiltViewModel
 import com.ramcosta.composedestinations.generated.app.destinations.E2EIEnrollmentScreenDestination
 import com.ramcosta.composedestinations.generated.app.destinations.HomeScreenDestination
 import com.ramcosta.composedestinations.generated.app.destinations.InitialSyncScreenDestination
 import com.ramcosta.composedestinations.generated.app.destinations.RemoveDeviceScreenDestination
 import com.wire.android.R
+import com.wire.android.di.wireViewModel
 import com.wire.android.navigation.BackStackMode
 import com.wire.android.navigation.NavigationCommand
 import com.wire.android.navigation.Navigator
@@ -86,7 +86,9 @@ import kotlinx.coroutines.launch
 fun LoginScreen(
     navigator: Navigator,
     loginNavArgs: LoginNavArgs,
-    loginEmailViewModel: LoginEmailViewModel = hiltViewModel()
+    loginEmailViewModel: LoginEmailViewModel = wireViewModel<LoginEmailViewModel, LoginEmailViewModel.Factory>(
+        creationCallback = { factory -> factory.create(loginNavArgs) }
+    )
 ) {
 
     LoginContent(
@@ -105,6 +107,7 @@ fun LoginScreen(
         onRemoveDeviceNeeded = {
             navigator.navigate(NavigationCommand(RemoveDeviceScreenDestination, BackStackMode.CLEAR_WHOLE))
         },
+        loginNavArgs = loginNavArgs,
         loginEmailViewModel = loginEmailViewModel,
         ssoLoginResult = loginNavArgs.ssoLoginResult,
         ssoCodeAutoLogin = loginNavArgs.ssoCodeAutoLogin
@@ -116,6 +119,7 @@ private fun LoginContent(
     onBackPressed: () -> Unit,
     onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     onRemoveDeviceNeeded: () -> Unit,
+    loginNavArgs: LoginNavArgs,
     loginEmailViewModel: LoginEmailViewModel,
     ssoLoginResult: DeepLinkResult.SSOLogin?,
     ssoCodeAutoLogin: SSOCodeAutoLogin?,
@@ -139,6 +143,7 @@ private fun LoginContent(
                     onBackPressed = onBackPressed,
                     onSuccess = onSuccess,
                     onRemoveDeviceNeeded = onRemoveDeviceNeeded,
+                    loginNavArgs = loginNavArgs,
                     loginEmailViewModel = loginEmailViewModel,
                     ssoLoginResult = ssoLoginResult,
                     ssoCodeAutoLogin = ssoCodeAutoLogin
@@ -154,6 +159,7 @@ private fun MainLoginContent(
     onBackPressed: () -> Unit,
     onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     onRemoveDeviceNeeded: () -> Unit,
+    loginNavArgs: LoginNavArgs,
     loginEmailViewModel: LoginEmailViewModel,
     ssoLoginResult: DeepLinkResult.SSOLogin?,
     ssoCodeAutoLogin: SSOCodeAutoLogin?,
@@ -233,6 +239,7 @@ private fun MainLoginContent(
                     LoginTabItem.SSO -> LoginSSOScreen(
                         onSuccess,
                         onRemoveDeviceNeeded,
+                        loginNavArgs,
                         ssoLoginResult,
                         ssoCodeAutoLogin,
                     )
@@ -264,7 +271,10 @@ private fun PreviewLoginScreen() = WireTheme {
             onBackPressed = {},
             onSuccess = { _, _ -> },
             onRemoveDeviceNeeded = {},
-            loginEmailViewModel = hiltViewModel(),
+            loginNavArgs = LoginNavArgs(),
+            loginEmailViewModel = wireViewModel<LoginEmailViewModel, LoginEmailViewModel.Factory>(
+                creationCallback = { factory -> factory.create(LoginNavArgs()) }
+            ),
             ssoLoginResult = null,
             ssoCodeAutoLogin = null
         )

app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginViewModel.kt

@@ -18,12 +18,9 @@
 
 package com.wire.android.ui.authentication.login
 
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.ClientScopeProvider
-import com.wire.android.di.KaliumCoreLogic
-import com.ramcosta.composedestinations.generated.app.navArgs
 import com.wire.kalium.logic.CoreLogic
 import com.wire.kalium.logic.configuration.server.ServerConfig
 import com.wire.kalium.logic.data.client.ClientCapability
@@ -32,37 +29,17 @@ import com.wire.kalium.logic.feature.auth.AddAuthenticatedUserUseCase
 import com.wire.kalium.logic.feature.auth.AuthenticationResult
 import com.wire.kalium.logic.feature.auth.DomainLookupUseCase
 import com.wire.kalium.logic.feature.client.RegisterClientResult
-import dagger.hilt.android.lifecycle.HiltViewModel
-import javax.inject.Inject
 
-@HiltViewModel
 @Suppress("TooManyFunctions")
 open class LoginViewModel(
-    savedStateHandle: SavedStateHandle,
+    loginNavArgs: LoginNavArgs,
     val clientScopeProviderFactory: ClientScopeProvider.Factory,
     val userDataStoreProvider: UserDataStoreProvider,
     val coreLogic: CoreLogic,
     private val loginExtension: LoginViewModelExtension,
     defaultServerConfig: ServerConfig.Links
 ) : ViewModel() {
 
-    @Inject
-    constructor(
-        savedStateHandle: SavedStateHandle,
-        clientScopeProviderFactory: ClientScopeProvider.Factory,
-        userDataStoreProvider: UserDataStoreProvider,
-        @KaliumCoreLogic coreLogic: CoreLogic,
-        defaultServerConfig: ServerConfig.Links
-    ) : this(
-        savedStateHandle,
-        clientScopeProviderFactory,
-        userDataStoreProvider,
-        coreLogic,
-        LoginViewModelExtension(clientScopeProviderFactory, userDataStoreProvider),
-        defaultServerConfig
-    )
-
-    private val loginNavArgs: LoginNavArgs = savedStateHandle.navArgs()
     val serverConfig: ServerConfig.Links = loginNavArgs.loginPasswordPath?.customServerConfig ?: defaultServerConfig
 
     suspend fun registerClient(

app/src/main/kotlin/com/wire/android/ui/authentication/login/SavedStateLoginSavedInputStore.kt

@@ -0,0 +1,52 @@
+/*
+ * Wire
+ * Copyright (C) 2026 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+
+package com.wire.android.ui.authentication.login
+
+import androidx.lifecycle.SavedStateHandle
+import dagger.Module
+import dagger.Provides
+import dagger.hilt.InstallIn
+import dagger.hilt.android.components.ViewModelComponent
+
+private const val USER_IDENTIFIER_SAVED_STATE_KEY = "user_identifier"
+private const val SSO_CODE_SAVED_STATE_KEY = "sso_code"
+
+class SavedStateLoginSavedInputStore(
+    private val savedStateHandle: SavedStateHandle,
+) : LoginSavedInputStore {
+    override var userIdentifier: String?
+        get() = savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY]
+        set(value) {
+            savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] = value
+        }
+
+    override var ssoCode: String?
+        get() = savedStateHandle[SSO_CODE_SAVED_STATE_KEY]
+        set(value) {
+            savedStateHandle[SSO_CODE_SAVED_STATE_KEY] = value
+        }
+}
+
+@Module
+@InstallIn(ViewModelComponent::class)
+object LoginSavedInputStoreModule {
+    @Provides
+    fun provideLoginSavedInputStore(savedStateHandle: SavedStateHandle): LoginSavedInputStore =
+        SavedStateLoginSavedInputStore(savedStateHandle)
+}

app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailVerificationCodeScreen.kt

@@ -20,7 +20,6 @@ package com.wire.android.ui.authentication.login.email
 
 import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.runtime.Composable
-import androidx.hilt.navigation.compose.hiltViewModel
 import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.authentication.verificationcode.VerificationCodeScreenContent
 import com.wire.android.ui.authentication.verificationcode.VerificationCodeState
@@ -29,7 +28,7 @@ import com.wire.android.util.ui.PreviewMultipleThemes
 
 @Composable
 fun LoginEmailVerificationCodeScreen(
-    viewModel: LoginEmailViewModel = hiltViewModel()
+    viewModel: LoginEmailViewModel
 ) = VerificationCodeScreenContent(
     viewModel.secondFactorVerificationCodeTextState,
     viewModel.secondFactorVerificationCodeState,

app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt

@@ -25,21 +25,21 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.DefaultWebSocketEnabledByDefault
 import com.wire.android.di.KaliumCoreLogic
 import com.wire.android.ui.authentication.login.LoginNavArgs
+import com.wire.android.ui.authentication.login.LoginSavedInputStore
 import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.authentication.login.LoginViewModel
+import com.wire.android.ui.authentication.login.LoginViewModelExtension
 import com.wire.android.ui.authentication.login.PreFilledUserIdentifierType
 import com.wire.android.ui.authentication.login.isProxyAuthRequired
 import com.wire.android.ui.authentication.login.toLoginError
 import com.wire.android.ui.authentication.verificationcode.VerificationCodeState
 import com.wire.android.ui.common.textfield.textAsFlow
-import com.ramcosta.composedestinations.generated.app.navArgs
 import com.wire.android.util.EMPTY
 import com.wire.android.util.dispatchers.DispatcherProvider
 import com.wire.android.util.ui.CountdownTimer
@@ -58,6 +58,9 @@
 import com.wire.kalium.logic.feature.auth.verification.RequestSecondFactorVerificationCodeUseCase
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import com.wire.kalium.logic.feature.session.CurrentSessionResult
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -68,28 +71,28 @@
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
-import javax.inject.Inject
 
 @Suppress("LongParameterList", "ComplexMethod", "TooManyFunctions")
-@HiltViewModel
-class LoginEmailViewModel @Inject constructor(
+@HiltViewModel(assistedFactory = LoginEmailViewModel.Factory::class)
+class LoginEmailViewModel @AssistedInject constructor(
+    @Assisted val loginNavArgs: LoginNavArgs,
     private val addAuthenticatedUser: AddAuthenticatedUserUseCase,
     clientScopeProviderFactory: ClientScopeProvider.Factory,
-    private val savedStateHandle: SavedStateHandle,
+    private val savedInputStore: LoginSavedInputStore,
     userDataStoreProvider: UserDataStoreProvider,
     @KaliumCoreLogic coreLogic: CoreLogic,
     private val resendCodeTimer: CountdownTimer,
     private val dispatchers: DispatcherProvider,
     defaultServerConfig: ServerConfig.Links,
     @DefaultWebSocketEnabledByDefault private val defaultWebSocketEnabledByDefault: Boolean,
 ) : LoginViewModel(
-    savedStateHandle,
+    loginNavArgs,
     clientScopeProviderFactory,
     userDataStoreProvider,
     coreLogic,
+    LoginViewModelExtension(clientScopeProviderFactory, userDataStoreProvider),
     defaultServerConfig
 ) {
-    val loginNavArgs: LoginNavArgs = savedStateHandle.navArgs()
     private val preFilledUserIdentifier: PreFilledUserIdentifierType = loginNavArgs.userHandle ?: PreFilledUserIdentifierType.None
 
     val userIdentifierTextState: TextFieldState = TextFieldState()
@@ -105,18 +108,23 @@
     @VisibleForTesting
     internal val loginJobData = MutableStateFlow<LoginJobData?>(null)
 
+    @AssistedFactory
+    interface Factory {
+        fun create(args: LoginNavArgs): LoginEmailViewModel
+    }
+
     init {
         userIdentifierTextState.setTextAndPlaceCursorAtEnd(
             if (preFilledUserIdentifier is PreFilledUserIdentifierType.PreFilled) {
                 preFilledUserIdentifier.userIdentifier
             } else {
-                savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] ?: String.EMPTY
+                savedInputStore.userIdentifier ?: String.EMPTY
             }
         )
         viewModelScope.launch {
             combine(
                 userIdentifierTextState.textAsFlow().distinctUntilChanged().onEach {
-                    savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] = it.toString()
+                    savedInputStore.userIdentifier = it.toString()
                 },
                 passwordTextState.textAsFlow(),
                 proxyIdentifierTextState.textAsFlow(),