Dagger + ViewModel + SavedStateHandle

ApplicationComponent.kt

@Singleton
@Component(modules = [AssistedInjectionModule::class])
interface ApplicationComponent {
    fun mySavedStateViewModelFactory(): MySavedStateViewModel.Factory
}

AssistedInjectionModule.java

@AssistedModule
@Module(includes = {AssistedInject_AssistedInjectionModule.class})
public class AssistedInjectionModule {
}

MySavedStateFragment.kt

import androidx.lifecycle.observe

class MySavedStateFragment : Fragment(R.layout.my_saved_state_fragment) {
    private val viewModel by navGraphSavedStateViewModels(R.id.registration_graph) { handle ->
        Injector.get().mySavedStateViewModelFactory().create(handle)
    }
  
    override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
        super.onViewCreated(view, savedInstanceState)
      
        val binding = MySavedStateFragmentBinding.bind(view)
      
        viewModel.blah.observe(viewLifecycleOwner) { blah ->
            // ...
        }
    }
}

MySavedStateViewModel.kt

class MySavedStateViewModel @AssistedInject constructor(
    private val authenticationManager: AuthenticationManager,
    @Assisted private val savedStateHandle: SavedStateHandle
) : ViewModel() {
    @AssistedInject.Factory
    interface Factory {
        fun create(savedStateHandle: SavedStateHandle): MySavedStateViewModel
    }
}

ViewModelUtils.kt

inline fun <reified T : ViewModel> SavedStateRegistryOwner.createAbstractSavedStateViewModelFactory(
    arguments: Bundle,
    crossinline creator: (SavedStateHandle) -> T
): ViewModelProvider.Factory {
    return object : AbstractSavedStateViewModelFactory(this, arguments) {
        @Suppress("UNCHECKED_CAST")
        override fun <T : ViewModel?> create(
            key: String, modelClass: Class<T>, handle: SavedStateHandle
        ): T = creator(handle) as T
    }
}

inline fun <reified T : ViewModel> Fragment.navGraphSavedStateViewModels(
    @IdRes navGraphId: Int,
    crossinline creator: (SavedStateHandle) -> T
): Lazy<T> {
    // Wrapped in lazy to not search the NavController each time we want the backStackEntry
    val backStackEntry by lazy { findNavController().getBackStackEntry(navGraphId) }

    return createViewModelLazy(T::class, storeProducer = {
        backStackEntry.viewModelStore
    }, factoryProducer = {
        backStackEntry.createAbstractSavedStateViewModelFactory(
            arguments = backStackEntry.arguments ?: Bundle(), creator = creator
        )
    })
}

inline fun <reified T : ViewModel> Fragment.fragmentSavedStateViewModels(
    crossinline creator: (SavedStateHandle) -> T
): Lazy<T> {
    return createViewModelLazy(T::class, storeProducer = {
        viewModelStore
    }, factoryProducer = {
        createAbstractSavedStateViewModelFactory(arguments ?: Bundle(), creator)
    })
}

I created a fork and added documentation to ViewModelUtils.kt as well as helper functions to manually inject CoroutineScope and a DispatcherProvider for Kotlin Coroutines in production and testing.