Blank Fragment shown on activity restart/config change in compose TabRow

I am trying to implement TabRow in compose with existing fragments being shown in FragmentContainerView. When the child fragments are added to the activity/host fragments they work fine. But upon activity restart the fragments dont appear to be visible but are present in the fragmentmanager but they don't get shown for some reason. Due to our existing implementation/architecture we only can add fragments and not replace them. Below is the sample code for it. There is no separate settings in Fragment A, B, C. They just contain one textView depicting their name.

Error Scenario gif available here https://i.sstatic.net/mxKiq.gif

Activity.kt

@Composable
fun MainScreen() {
        var selectedTabIndex by remember { mutableStateOf(0) }
        Column(Modifier.fillMaxSize()) {
            TabRow(
                selectedTabIndex = selectedTabIndex,
                backgroundColor = Color.Gray,
            ) {
                tabs.forEachIndexed { index, title ->
                    Tab(
                        text = {
                            Text(
                                text = title,
                                maxLines = 1,
                                overflow = TextOverflow.Ellipsis
                            )
                        },
                        selected = selectedTabIndex == index,
                        onClick = {
                            selectedTabIndex = index
                        },
                        selectedContentColor = Color.Black,
                        unselectedContentColor = Color.White
                    )
                }
            }
            FragmentContainer(
                modifier = Modifier.fillMaxSize(),
                commit = getCommitFunction(
                    index = selectedTabIndex
                )
            )
        }
    }

    private fun getFragment(index: Int): Fragment {
        return when (index) {
            0 -> FragmentA.newInstance()
            1 -> FragmentB.newInstance()
            else -> FragmentC.newInstance()
        }
    }

    private fun getCommitFunction(
        index: Int,
    ): FragmentTransaction.(containerId: Int) -> Unit =
        { containerId ->
            showPage(this, index, containerId)
        }

    private fun showPage(transaction: FragmentTransaction, position: Int, containerId: Int) {

        // hide current showing fragment
        tabs.forEach { fragmentTag ->
            val fragment: Fragment? = supportFragmentManager.findFragmentByTag(fragmentTag)
            if (fragment != null && fragment.isVisible) {
                transaction.hide(fragment)
            }
        }

        val fragmentTag = tabs[position]
        var fragment: Fragment? = supportFragmentManager.findFragmentByTag(fragmentTag)

        if (fragment == null) {
            // we don't have the fragment yet, need to create it
            fragment = getFragment(position)
            println("containerId $containerId")
            transaction.add(containerId, fragment, fragmentTag)
        } else {
            println("containerId else $containerId")
            // we already created the fragment, just need to show it
            transaction.show(fragment)
        }

        transaction.setPrimaryNavigationFragment(fragment)
    }

FragmentContainer.kt

@Composable
fun FragmentContainer(
    modifier: Modifier = Modifier,
    commit: FragmentTransaction.(containerId: Int) -> Unit
) {
    val localView = LocalView.current
    // Find the parent fragment, if one exists. This will let us ensure that
    // fragments inflated via a FragmentContainerView are properly nested
    // (which, in turn, allows the fragments to properly save/restore their state)
    val parentFragment = remember(localView) {
        try {
            localView.findFragment<Fragment>()
        } catch (e: IllegalStateException) {
            // findFragment throws if no parent fragment is found
            null
        }
    }
    val containerId by rememberSaveable { mutableStateOf(View.generateViewId()) }
    val container = remember { mutableStateOf<FragmentContainerView?>(null) }
    val viewBlock: (Context) -> View = remember(localView) {
        { context ->
            FragmentContainerView(context).apply { id = containerId }
        }
    }
    AndroidView(
        modifier = modifier,
        factory = viewBlock,
        update = {view ->
            val fragmentManager = parentFragment?.childFragmentManager
                ?: (view.context as? FragmentActivity)?.supportFragmentManager
            fragmentManager?.commit { commit(view.id) }
            container.value = view as FragmentContainerView
        }
    )

    // Set up a DisposableEffect that will clean up fragments when the FragmentContainer is disposed
    val localContext = LocalContext.current
    DisposableEffect(localView, localContext, container) {
        onDispose {
            val fragmentManager = parentFragment?.childFragmentManager
                ?: (localContext as? FragmentActivity)?.supportFragmentManager
            // Now find the fragment inflated via the FragmentContainerView
            val existingFragment = fragmentManager?.findFragmentById(container.value?.id ?: 0)
            if (existingFragment != null && !fragmentManager.isStateSaved) {
                // If the state isn't saved, that means that some state change
                // has removed this Composable from the hierarchy
                fragmentManager.commit(true) {
                    remove(existingFragment)
                }
            }
        }
    }
}

The reason the state is not restored is because nothing is storing the state of this. It gets recreated every time the app has been killed and restored.

        FragmentContainer(
            modifier = Modifier.fillMaxSize(),
            commit = getCommitFunction(
                index = selectedTabIndex
            )
        )

We need something to keep the state of these fragments. One way is to use NavHost and manage the navigation through Navigator.

I've created a simple example in https://github.com/elye/demo_android_jetpack_compose_fragment_navigation. Refer to the last example of the design on Tab. Hope this helps.