Example on how to to inject SavedStateHandle into ViewModel's

Example.kt

// ViewModel factory

@MainThread
inline fun <reified VM : ViewModel> ComponentActivity.viewModel(
    noinline provider: (SavedStateHandle) -> VM
) = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { this },
    viewModelStoreOwnerProducer = { this },
    viewModelProvider = provider
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.viewModel(
    noinline provider: (SavedStateHandle) -> VM
) = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { this },
    viewModelStoreOwnerProducer = { this },
    viewModelProvider = provider
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.parentViewModel() = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { requireParentFragment() },
    viewModelStoreOwnerProducer = { requireParentFragment() },
    viewModelProvider = {
        throw IllegalStateException("${VM::class.java.name} should already exist scoped to the parent fragment.")
    }
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.activityViewModel() = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { requireActivity() },
    viewModelStoreOwnerProducer = { requireActivity() },
    viewModelProvider = {
        throw IllegalStateException("${VM::class.java.name} should already exist scoped to the activity.")
    }
)

fun <VM : ViewModel> createLazyViewModel(
    viewModelClass: KClass<VM>,
    savedStateRegistryOwnerProducer: () -> SavedStateRegistryOwner,
    viewModelStoreOwnerProducer: () -> ViewModelStoreOwner,
    viewModelProvider: (SavedStateHandle) -> VM
) = ViewModelLazy(viewModelClass, { viewModelStoreOwnerProducer().viewModelStore }) {
    object : AbstractSavedStateViewModelFactory(savedStateRegistryOwnerProducer(), Bundle()) {
        @Suppress("UNCHECKED_CAST")
        override fun <T : ViewModel?> create(
            key: String,
            modelClass: Class<T>,
            handle: SavedStateHandle
        ): T {
            return viewModelProvider(handle) as T
        }
    }
}


// ViewModel

class EditGoalNameViewModel @AssistedInject constructor(
    @Assisted args: EditGoalNameScreen.Args,
    @Assisted savedStateHandle: SavedStateHandle
) : BaseViewModel() {

    @AssistedInject.Factory
    interface Factory {
        fun create(
            args: EditGoalNameScreen.Args,
            savedStateHandle: SavedStateHandle
        ): EditGoalNameViewModel
    }
  
    ...
}


// Fragment

class EditGoalNameFragment @Inject constructor(
    viewModelFactory: EditGoalNameViewModel.Factory
) : BaseFragment<EditGoalNameScreen>(R.layout.fragment_edit_goal_name) {

    private val viewModel by viewModel { viewModelFactory.create(screen.args, it) }
  
    ...
}