diff --git a/src/com/android/documentsui/BaseActivity.java b/src/com/android/documentsui/BaseActivity.java
index 4228ead..9e38303 100644
--- a/src/com/android/documentsui/BaseActivity.java
+++ b/src/com/android/documentsui/BaseActivity.java
@@ -54,6 +54,7 @@
 import com.android.documentsui.Injector.Injected;
 import com.android.documentsui.NavigationViewManager.Breadcrumb;
 import com.android.documentsui.base.DocumentInfo;
+import com.android.documentsui.base.DocumentStack;
 import com.android.documentsui.base.EventHandler;
 import com.android.documentsui.base.RootInfo;
 import com.android.documentsui.base.Shared;
@@ -87,6 +88,7 @@
         extends AppCompatActivity implements CommonAddons, NavigationViewManager.Environment {
 
     private static final String BENCHMARK_TESTING_PACKAGE = "com.android.documentsui.appperftests";
+    private static final String TAG = "BaseActivity";
 
     protected SearchViewManager mSearchManager;
     protected AppsRowManager mAppsRowManager;
@@ -118,10 +120,25 @@
 
     private PreferencesMonitor mPreferencesMonitor;
 
-    private boolean mHasProfileBecomeUnavailable = false;
+    private final DocumentStack mInitialStack = new DocumentStack();
+    private UserId mLastSelectedUser = null;
 
-    public void setHasProfileBecomeUnavailable(boolean hasProfileBecomeUnavailable) {
-        mHasProfileBecomeUnavailable = hasProfileBecomeUnavailable;
+    protected void setInitialStack(DocumentStack stack) {
+        if (mInitialStack.isInitialized()) {
+            if (DEBUG) {
+                Log.d(TAG, "Initial stack already initialised. " + mInitialStack.isInitialized());
+            }
+            return;
+        }
+        mInitialStack.reset(stack);
+    }
+
+    public DocumentStack getInitialStack() {
+        return mInitialStack;
+    }
+
+    public UserId getLastSelectedUser() {
+        return mLastSelectedUser;
     }
 
     public BaseActivity(@LayoutRes int layoutId, String tag) {
@@ -140,7 +157,7 @@
 
     @VisibleForTesting
     protected void initConfigStore() {
-        mConfigStore = DocumentsApplication.getConfigStore(this);
+        mConfigStore = DocumentsApplication.getConfigStore();
     }
 
     @VisibleForTesting
@@ -287,6 +304,12 @@
 
         mUserIdManager = DocumentsApplication.getUserIdManager(this);
         mUserManagerState = DocumentsApplication.getUserManagerState(this);
+        // If private space feature flag is enabled, we should store the intent that launched docsUi
+        // so that we can use this intent to get CrossProfileResolveInfo when ever we want to,
+        // for example when ACTION_PROFILE_AVAILABLE intent is received
+        if (mUserManagerState != null) {
+            mUserManagerState.setCurrentStateIntent(intent);
+        }
         mSearchManager = new SearchViewManager(searchListener, queryInterceptor,
                 chipGroup, savedInstanceState);
         // initialize the chip sets by accept mime types
@@ -339,9 +362,7 @@
                 // When a profile with user property SHOW_IN_QUIET_MODE_HIDDEN is currently
                 // selected, and it becomes unavailable, we reset the roots to recents.
                 // We do not reset it to recents when pick activity is due to ACTION_CREATE_DOCUMENT
-                mInjector.actions.loadCrossProfileRoot(
-                        (mHasProfileBecomeUnavailable && mState.action != State.ACTION_CREATE)
-                                ? getRecentsRoot() : getCurrentRoot(), userId);
+                mInjector.actions.loadCrossProfileRoot(getCurrentRoot(), userId);
             }
         });
 
@@ -392,6 +413,12 @@
     }
 
     @Override
+    protected void onPause() {
+        super.onPause();
+        mLastSelectedUser = getSelectedUser();
+    }
+
+    @Override
     public boolean onCreateOptionsMenu(Menu menu) {
         boolean showMenu = super.onCreateOptionsMenu(menu);
 
@@ -896,10 +923,6 @@
         }
     }
 
-    public RootInfo getRecentsRoot() {
-        return mProviders.generateRecentsRoot(getSelectedUser());
-    }
-
     @Override
     public DocumentInfo getCurrentDirectory() {
         return mState.stack.peek();
diff --git a/src/com/android/documentsui/DocumentsApplication.java b/src/com/android/documentsui/DocumentsApplication.java
index 92aa3f7..84e4754 100644
--- a/src/com/android/documentsui/DocumentsApplication.java
+++ b/src/com/android/documentsui/DocumentsApplication.java
@@ -83,20 +83,7 @@
     private Lookup<String, String> mFileTypeLookup;
 
     public static ProvidersCache getProvidersCache(Context context) {
-        ProvidersCache providers =
-                ((DocumentsApplication) context.getApplicationContext()).mProviders;
-        final ConfigStore configStore = getConfigStore(context);
-        // When private space in DocsUI is enabled then ProvidersCache should use UserManagerState
-        // else it should use UserIdManager. The following if-check will ensure the construction of
-        // a new ProvidersCache instance whenever there is a mismatch in this.
-        if (configStore.isPrivateSpaceInDocsUIEnabled()
-                != providers.isProvidersCacheUsingUserManagerState()) {
-            providers = configStore.isPrivateSpaceInDocsUIEnabled()
-                    ? new ProvidersCache(context, getUserManagerState(context), configStore)
-                    : new ProvidersCache(context, getUserIdManager(context), configStore);
-            ((DocumentsApplication) context.getApplicationContext()).mProviders = providers;
-        }
-        return providers;
+        return ((DocumentsApplication) context.getApplicationContext()).mProviders;
     }
 
     public static ThumbnailCache getThumbnailCache(Context context) {
@@ -140,7 +127,7 @@
     public static UserManagerState getUserManagerState(Context context) {
         UserManagerState userManagerState =
                 ((DocumentsApplication) context.getApplicationContext()).mUserManagerState;
-        if (userManagerState == null && getConfigStore(context).isPrivateSpaceInDocsUIEnabled()) {
+        if (userManagerState == null && getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
             userManagerState = UserManagerState.create(context);
             ((DocumentsApplication) context.getApplicationContext()).mUserManagerState =
                     userManagerState;
@@ -159,7 +146,7 @@
     /**
      * Retrieve {@link ConfigStore} instance to access feature flags in production code.
      */
-    public static synchronized ConfigStore getConfigStore(Context context) {
+    public static synchronized ConfigStore getConfigStore() {
         if (sConfigStore == null) {
             sConfigStore = new ConfigStore.ConfigStoreImpl();
         }
@@ -167,14 +154,6 @@
     }
 
     /**
-     * Set {@link #mProviders} as null onDestroy of BaseActivity so that new session uses new
-     * instance of {@link #mProviders}
-     */
-    public static void invalidateProvidersCache(Context context) {
-        ((DocumentsApplication) context.getApplicationContext()).mProviders = null;
-    }
-
-    /**
      * Set {@link #mUserManagerState} as null onDestroy of BaseActivity so that new session uses new
      * instance of {@link #mUserManagerState}
      */
@@ -217,19 +196,14 @@
             Log.w(TAG, "Can't obtain OverlayManager from System Service!");
         }
 
-        if (getConfigStore(this).isPrivateSpaceInDocsUIEnabled()) {
+        if (getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
             mUserManagerState = UserManagerState.create(this);
             mUserIdManager = null;
-            synchronized (DocumentsApplication.class) {
-                mProviders = new ProvidersCache(this, mUserManagerState, getConfigStore(this));
-            }
         } else {
             mUserManagerState = null;
             mUserIdManager = UserIdManager.create(this);
-            synchronized (DocumentsApplication.class) {
-                mProviders = new ProvidersCache(this, mUserIdManager, getConfigStore(this));
-            }
         }
+        mProviders = new ProvidersCache(this);
 
         mProviders.updateAsync(/* forceRefreshAll= */ false, /* callback= */  null);
 
@@ -288,13 +262,15 @@
                 final String packageName = data.getSchemeSpecificPart();
                 mProviders.updatePackageAsync(UserId.DEFAULT_USER, packageName);
             } else if (PROFILE_FILTER_ACTIONS.contains(action)) {
-                // After we have reloaded roots. Resend the broadcast locally so the other
-                // components can reload properly after roots are updated.
-                if (getConfigStore(context).isPrivateSpaceInDocsUIEnabled()) {
+                // Make the changes to UserManagerState object before calling providers updateAsync
+                // so that providers for all the users are loaded
+                if (getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
                     UserHandle userHandle = intent.getParcelableExtra(Intent.EXTRA_USER);
                     UserId userId = UserId.of(userHandle);
                     getUserManagerState(context).onProfileActionStatusChange(action, userId);
                 }
+                // After we have reloaded roots. Resend the broadcast locally so the other
+                // components can reload properly after roots are updated.
                 mProviders.updateAsync(/* forceRefreshAll= */ true,
                         () -> LocalBroadcastManager.getInstance(context).sendBroadcast(intent));
             } else {
diff --git a/src/com/android/documentsui/MultiRootDocumentsLoader.java b/src/com/android/documentsui/MultiRootDocumentsLoader.java
index f450814..db78daa 100644
--- a/src/com/android/documentsui/MultiRootDocumentsLoader.java
+++ b/src/com/android/documentsui/MultiRootDocumentsLoader.java
@@ -59,7 +59,7 @@
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
-/*
+/**
  * The abstract class to query multiple roots from {@link android.provider.DocumentsProvider}
  * and return the combined result.
  */
@@ -88,7 +88,7 @@
     private LockingContentObserver mObserver;
 
     @GuardedBy("mTasks")
-    /** A authority -> QueryTask map */
+    /* A authority -> QueryTask map */
     private final Map<String, QueryTask> mTasks = new HashMap<>();
 
     private CountDownLatch mFirstPassLatch;
@@ -96,7 +96,7 @@
 
     private DirectoryResult mResult;
 
-    /*
+    /**
      * Create the loader to query roots from {@link android.provider.DocumentsProvider}.
      *
      * @param context the context
diff --git a/src/com/android/documentsui/NavigationViewManager.java b/src/com/android/documentsui/NavigationViewManager.java
index 6727f1c..c376c86 100644
--- a/src/com/android/documentsui/NavigationViewManager.java
+++ b/src/com/android/documentsui/NavigationViewManager.java
@@ -71,7 +71,8 @@
     private final ConfigStore mConfigStore;
 
     private boolean mIsActionModeActivated = false;
-    private @ColorRes int mDefaultStatusBarColorResId;
+    @ColorRes
+    private int mDefaultStatusBarColorResId;
 
     public NavigationViewManager(
             BaseActivity activity,
diff --git a/src/com/android/documentsui/ProfileTabs.java b/src/com/android/documentsui/ProfileTabs.java
index d096d59..b692c7a 100644
--- a/src/com/android/documentsui/ProfileTabs.java
+++ b/src/com/android/documentsui/ProfileTabs.java
@@ -128,13 +128,8 @@
             // Update the layout according to the current root if necessary.
             // Make sure we do not invoke callback. Otherwise, it is likely to cause infinite loop.
             mTabs.removeOnTabSelectedListener(mOnTabSelectedListener);
-            if (!mUserIds.contains(currentRoot.userId)) {
-                mTabs.addOnTabSelectedListener(mOnTabSelectedListener);
-                mTabs.selectTab(mTabs.getTabAt(mUserIds.indexOf(UserId.CURRENT_USER)));
-            } else {
-                mTabs.selectTab(mTabs.getTabAt(mUserIds.indexOf(currentRoot.userId)));
-                mTabs.addOnTabSelectedListener(mOnTabSelectedListener);
-            }
+            mTabs.selectTab(mTabs.getTabAt(mUserIds.indexOf(currentRoot.userId)));
+            mTabs.addOnTabSelectedListener(mOnTabSelectedListener);
         }
         mTabsContainer.setVisibility(shouldShow() ? View.VISIBLE : View.GONE);
 
diff --git a/src/com/android/documentsui/UserManagerState.java b/src/com/android/documentsui/UserManagerState.java
index 28a83e8..eccc6b0 100644
--- a/src/com/android/documentsui/UserManagerState.java
+++ b/src/com/android/documentsui/UserManagerState.java
@@ -95,6 +95,11 @@
     void onProfileActionStatusChange(String action, UserId userId);
 
     /**
+     * Sets the intent that triggered the launch of the DocsUI
+     */
+    void setCurrentStateIntent(Intent intent);
+
+    /**
      * Creates an implementation of {@link UserManagerState}.
      */
     // TODO: b/314746383 Make this class a singleton
@@ -136,6 +141,7 @@
         @GuardedBy("mCanFrowardToProfileIdMap")
         private final Map<UserId, Boolean> mCanFrowardToProfileIdMap = new HashMap<>();
 
+        private Intent mCurrentStateIntent;
 
         private final BroadcastReceiver mIntentReceiver = new BroadcastReceiver() {
             @Override
@@ -159,7 +165,7 @@
         private RuntimeUserManagerState(Context context) {
             this(context, UserId.CURRENT_USER,
                     Features.CROSS_PROFILE_TABS && isDeviceSupported(context),
-                    DocumentsApplication.getConfigStore(context));
+                    DocumentsApplication.getConfigStore());
         }
 
         @VisibleForTesting
@@ -185,7 +191,6 @@
         public List<UserId> getUserIds() {
             synchronized (mUserIds) {
                 if (mUserIds.isEmpty()) {
-                    Log.d("profileAction", "user ids empty");
                     mUserIds.addAll(getUserIdsInternal());
                 }
                 return mUserIds;
@@ -234,15 +239,6 @@
                 synchronized (mUserIds) {
                     mUserIds.remove(userId);
                 }
-                synchronized (mUserIdToLabelMap) {
-                    mUserIdToLabelMap.remove(userId);
-                }
-                synchronized (mUserIdToBadgeMap) {
-                    mUserIdToBadgeMap.remove(userId);
-                }
-                synchronized (mCanFrowardToProfileIdMap) {
-                    mCanFrowardToProfileIdMap.remove(userId);
-                }
             } else if (Intent.ACTION_PROFILE_AVAILABLE.equals(action)) {
                 synchronized (mUserIds) {
                     if (!mUserIds.contains(userId)) {
@@ -250,19 +246,40 @@
                     }
                 }
                 synchronized (mUserIdToLabelMap) {
-                    mUserIdToLabelMap.put(userId, getProfileLabel(userId));
+                    if (!mUserIdToLabelMap.containsKey(userId)) {
+                        mUserIdToLabelMap.put(userId, getProfileLabel(userId));
+                    }
                 }
                 synchronized (mUserIdToBadgeMap) {
-                    mUserIdToBadgeMap.put(userId, getProfileBadge(userId));
+                    if (!mUserIdToBadgeMap.containsKey(userId)) {
+                        mUserIdToBadgeMap.put(userId, getProfileBadge(userId));
+                    }
                 }
                 synchronized (mCanFrowardToProfileIdMap) {
-                    mCanFrowardToProfileIdMap.put(userId, true);
+                    if (!mCanFrowardToProfileIdMap.containsKey(userId)) {
+                        if (userId.getIdentifier() == ActivityManager.getCurrentUser()
+                                || isCrossProfileContentSharingStrategyDelegatedFromParent(
+                                UserHandle.of(userId.getIdentifier()))
+                                || CrossProfileUtils.getCrossProfileResolveInfo(mCurrentUser,
+                                mContext.getPackageManager(), mCurrentStateIntent, mContext,
+                                mConfigStore.isPrivateSpaceInDocsUIEnabled()) != null) {
+                            mCanFrowardToProfileIdMap.put(userId, true);
+                        } else {
+                            mCanFrowardToProfileIdMap.put(userId, false);
+                        }
+
+                    }
                 }
             } else {
                 Log.e(TAG, "Unexpected action received: " + action);
             }
         }
 
+        @Override
+        public void setCurrentStateIntent(Intent intent) {
+            mCurrentStateIntent = intent;
+        }
+
         private List<UserId> getUserIdsInternal() {
             final List<UserId> result = new ArrayList<>();
 
@@ -538,8 +555,7 @@
             /*
              * Cross profile resolve info need to be checked in the following 2 cases:
              * 1. current user is either parent or delegates check to parent and the target user
-             * does
-             *    not delegate to parent
+             *    does not delegate to parent
              * 2. current user does not delegate check to the parent and the target user is the
              *    parent profile
              */
diff --git a/src/com/android/documentsui/dirlist/DirectoryFragment.java b/src/com/android/documentsui/dirlist/DirectoryFragment.java
index 32f3683..bf54e25 100644
--- a/src/com/android/documentsui/dirlist/DirectoryFragment.java
+++ b/src/com/android/documentsui/dirlist/DirectoryFragment.java
@@ -246,7 +246,6 @@
             } else {
                 profileStatusReceiverPreV(intent, action);
             }
-
         }
 
         private void profileStatusReceiverPostV(Intent intent, String action) {
@@ -310,16 +309,17 @@
 
     private void onHiddenProfileStatusChange(String action, UserId userId) {
         if (Intent.ACTION_PROFILE_UNAVAILABLE.equals(action)) {
-            mActivity.setHasProfileBecomeUnavailable(true);
             if (mProviderTestRunnable != null) {
                 mHandler.removeCallbacks(mProviderTestRunnable);
                 mProviderTestRunnable = null;
             }
-            while (mState.stack.size() > 0) {
-                mState.stack.pop();
+            if (!mActivity.isSearchExpanded()) {
+                if (mActivity.getLastSelectedUser() != null
+                        && mActivity.getLastSelectedUser().equals(userId)) {
+                    mState.stack.reset(mActivity.getInitialStack());
+                }
+                mActivity.refreshCurrentRootAndDirectory(AnimationView.ANIM_NONE);
             }
-            mActivity.updateNavigator();
-            mActivity.setHasProfileBecomeUnavailable(false);
         } else {
             checkUriAndScheduleCheckIfNeeded(userId);
         }
diff --git a/src/com/android/documentsui/files/FilesActivity.java b/src/com/android/documentsui/files/FilesActivity.java
index 053e5db..7a53e53 100644
--- a/src/com/android/documentsui/files/FilesActivity.java
+++ b/src/com/android/documentsui/files/FilesActivity.java
@@ -300,7 +300,7 @@
     }
 
     @Override
-    public void onDestroy() {
+    protected void onDestroy() {
         super.onDestroy();
     }
 
@@ -352,6 +352,8 @@
         final RootInfo root = getCurrentRoot();
         final DocumentInfo cwd = getCurrentDirectory();
 
+        setInitialStack(mState.stack);
+
         assert (!mSearchManager.isSearching());
 
         if (mState.stack.isRecents()) {
diff --git a/src/com/android/documentsui/picker/PickActivity.java b/src/com/android/documentsui/picker/PickActivity.java
index 2061a8f..c14c8b8 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -300,6 +300,11 @@
     }
 
     @Override
+    protected void onDestroy() {
+        super.onDestroy();
+    }
+
+    @Override
     public String getDrawerTitle() {
         String title;
         try {
@@ -353,6 +358,8 @@
         final RootInfo root = getCurrentRoot();
         final DocumentInfo cwd = getCurrentDirectory();
 
+        setInitialStack(mState.stack);
+
         if (mState.stack.isRecents()) {
             DirectoryFragment.showRecentsOpen(fm, anim);
 
diff --git a/src/com/android/documentsui/roots/ProvidersCache.java b/src/com/android/documentsui/roots/ProvidersCache.java
index 080b729..15089a6 100644
--- a/src/com/android/documentsui/roots/ProvidersCache.java
+++ b/src/com/android/documentsui/roots/ProvidersCache.java
@@ -49,11 +49,8 @@
 import androidx.annotation.Nullable;
 import androidx.localbroadcastmanager.content.LocalBroadcastManager;
 
-import com.android.documentsui.ConfigStore;
 import com.android.documentsui.DocumentsApplication;
 import com.android.documentsui.R;
-import com.android.documentsui.UserIdManager;
-import com.android.documentsui.UserManagerState;
 import com.android.documentsui.UserPackage;
 import com.android.documentsui.archives.ArchivesProvider;
 import com.android.documentsui.base.LookupApplicationName;
@@ -123,33 +120,14 @@
     @GuardedBy("mObservedAuthoritiesDetails")
     private final Map<UserAuthority, PackageDetails> mObservedAuthoritiesDetails = new HashMap<>();
 
-    private final UserIdManager mUserIdManager;
-    private final UserManagerState mUserManagerState;
-    private final ConfigStore mConfigStore;
-
-    public ProvidersCache(Context context, UserIdManager userIdManager, ConfigStore configStore) {
+    public ProvidersCache(Context context) {
         mContext = context;
-        mUserIdManager = userIdManager;
-        mUserManagerState = null;
-        mConfigStore = configStore;
-    }
-
-    public ProvidersCache(Context context, UserManagerState userManagerState,
-            ConfigStore configStore) {
-        mContext = context;
-        mUserIdManager = null;
-        mUserManagerState = userManagerState;
-        mConfigStore = configStore;
-    }
-
-    public boolean isProvidersCacheUsingUserManagerState() {
-        return mUserManagerState != null;
     }
 
     /**
      * Generates recent root for the provided user id
      */
-    public RootInfo generateRecentsRoot(UserId rootUserId) {
+    private RootInfo generateRecentsRoot(UserId rootUserId) {
         return new RootInfo() {{
             // Special root for recents
             userId = rootUserId;
@@ -220,7 +198,7 @@
         // For that reason we update our RecentsRoot to reflect
         // the current language.
         final String title = mContext.getString(R.string.root_recent);
-        List<UserId> userIds = getUserIds();
+        List<UserId> userIds = new ArrayList<>(getUserIds());
         for (UserId userId : userIds) {
             RootInfo recentRoot = createOrGetRecentsRoot(userId);
             recentRoot.title = title;
@@ -559,7 +537,7 @@
 
             final long start = SystemClock.elapsedRealtime();
 
-            List<UserId> userIds = getUserIds();
+            List<UserId> userIds = new ArrayList<>(getUserIds());
             for (UserId userId : userIds) {
                 final RootInfo recents = createOrGetRecentsRoot(userId);
                 synchronized (mLock) {
@@ -727,9 +705,9 @@
     }
 
     private List<UserId> getUserIds() {
-        if (mConfigStore.isPrivateSpaceInDocsUIEnabled()) {
-            return mUserManagerState.getUserIds();
+        if (DocumentsApplication.getConfigStore().isPrivateSpaceInDocsUIEnabled()) {
+            return DocumentsApplication.getUserManagerState(mContext).getUserIds();
         }
-        return mUserIdManager.getUserIds();
+        return DocumentsApplication.getUserIdManager(mContext).getUserIds();
     }
 }
diff --git a/tests/common/com/android/documentsui/TestUserManagerState.java b/tests/common/com/android/documentsui/TestUserManagerState.java
index 8dbe62a..5f41ce3 100644
--- a/tests/common/com/android/documentsui/TestUserManagerState.java
+++ b/tests/common/com/android/documentsui/TestUserManagerState.java
@@ -18,6 +18,7 @@
 
 import android.content.Intent;
 import android.graphics.drawable.Drawable;
+import android.util.Log;
 
 import com.android.documentsui.base.UserId;
 
@@ -28,12 +29,14 @@
 
 public class TestUserManagerState implements UserManagerState {
 
+    private static final String TAG = "TestUserManagerState";
     public List<UserId> userIds = new ArrayList<>();
     public Map<UserId, String> userIdToLabelMap = new HashMap<>();
     public Map<UserId, Boolean> canFrowardToProfileIdMap = new HashMap<>();
     public Map<UserId, Drawable> userIdToBadgeMap = new HashMap<>();
     public String profileLabel = "Test";
     public Drawable profileBadge = null;
+    public Intent intent = new Intent();
 
     @Override
     public List<UserId> getUserIds() {
@@ -59,16 +62,26 @@
     public void onProfileActionStatusChange(String action, UserId userId) {
         if (Intent.ACTION_PROFILE_UNAVAILABLE.equals(action)) {
             userIds.remove(userId);
-            userIdToLabelMap.remove(userId);
-            userIdToBadgeMap.remove(userId);
-            canFrowardToProfileIdMap.put(userId, false);
             return;
+        } else if (Intent.ACTION_PROFILE_AVAILABLE.equals(action)) {
+            if (!userIds.contains(userId)) {
+                userIds.add(userId);
+            }
+            if (!userIdToLabelMap.containsKey(userId)) {
+                userIdToLabelMap.put(userId, profileLabel);
+            }
+            if (!userIdToBadgeMap.containsKey(userId)) {
+                userIdToBadgeMap.put(userId, profileBadge);
+            }
+            if (!canFrowardToProfileIdMap.containsKey(userId)) {
+                canFrowardToProfileIdMap.put(userId, true);
+            }
+        } else {
+            Log.e(TAG, "Unexpected action received: " + action);
         }
-        if (!userIds.contains(userId)) {
-            userIds.add(userId);
-        }
-        userIdToLabelMap.put(userId, profileLabel);
-        userIdToBadgeMap.put(userId, profileBadge);
-        canFrowardToProfileIdMap.put(userId, true);
+    }
+
+    @Override
+    public void setCurrentStateIntent(Intent intent) {
     }
 }
diff --git a/tests/unit/com/android/documentsui/UserManagerStateTest.java b/tests/unit/com/android/documentsui/UserManagerStateTest.java
index 25c7607..afe8f45 100644
--- a/tests/unit/com/android/documentsui/UserManagerStateTest.java
+++ b/tests/unit/com/android/documentsui/UserManagerStateTest.java
@@ -16,22 +16,33 @@
 
 package com.android.documentsui;
 
+import static com.android.documentsui.DevicePolicyResources.Drawables.Style.SOLID_COLORED;
+import static com.android.documentsui.DevicePolicyResources.Drawables.WORK_PROFILE_ICON;
+import static com.android.documentsui.DevicePolicyResources.Strings.PERSONAL_TAB;
+import static com.android.documentsui.DevicePolicyResources.Strings.WORK_TAB;
+
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
 import android.Manifest;
+import android.app.admin.DevicePolicyManager;
+import android.app.admin.DevicePolicyResourcesManager;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
 import android.content.pm.UserProperties;
+import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
 import android.os.UserManager;
 
 import androidx.test.filters.SdkSuppress;
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.documentsui.base.UserId;
 import com.android.documentsui.testing.UserManagers;
@@ -52,6 +63,10 @@
 @SdkSuppress(minSdkVersion = 31, codeName = "S")
 public class UserManagerStateTest {
 
+    private static final String PERSONAL = "Personal";
+    private static final String WORK = "Work";
+    private static final String PRIVATE = "Private";
+
     private final UserHandle mSystemUser = UserHandle.SYSTEM;
     private final UserHandle mManagedUser = UserHandle.of(100);
     private final UserHandle mPrivateUser = UserHandle.of(101);
@@ -66,6 +81,7 @@
     private final Intent mMockIntent = mock(Intent.class);
     private final UserManager mMockUserManager = UserManagers.create();
     private final PackageManager mMockPackageManager = mock(PackageManager.class);
+    private final DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
     private UserManagerState mUserManagerState;
 
     @Before
@@ -114,7 +130,7 @@
             when(mMockUserManager.getUserProperties(mNormalUser)).thenReturn(normalUserProperties);
         }
 
-        when(mMockUserManager.getProfileParent(mSystemUser)).thenReturn(mSystemUser);
+        when(mMockUserManager.getProfileParent(mSystemUser)).thenReturn(null);
         when(mMockUserManager.getProfileParent(mManagedUser)).thenReturn(mSystemUser);
         when(mMockUserManager.getProfileParent(mPrivateUser)).thenReturn(mSystemUser);
         when(mMockUserManager.getProfileParent(mOtherUser)).thenReturn(mSystemUser);
@@ -129,6 +145,12 @@
         when(mMockContext.getPackageManager()).thenReturn(mMockPackageManager);
         when(mMockContext.getSystemServiceName(UserManager.class)).thenReturn("mMockUserManager");
         when(mMockContext.getSystemService(UserManager.class)).thenReturn(mMockUserManager);
+        when(mMockContext.getSystemServiceName(DevicePolicyManager.class))
+                .thenReturn(Context.DEVICE_POLICY_SERVICE);
+        when(mMockContext.getSystemService(Context.DEVICE_POLICY_SERVICE))
+                .thenReturn(mDevicePolicyManager);
+        when(mMockContext.getResources()).thenReturn(
+                InstrumentationRegistry.getInstrumentation().getTargetContext().getResources());
     }
 
     @Test
@@ -560,7 +582,7 @@
     }
 
     @Test
-    public void testOnProfileStatusChange_anyIntentActionOnManagedProfile() {
+    public void testOnProfileStatusChange_anyIntentActionForManagedProfile() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         initializeUserManagerState(currentUser,
@@ -584,7 +606,7 @@
     }
 
     @Test
-    public void testOnProfileStatusChange_actionProfileUnavailableOnPrivateProfile() {
+    public void testOnProfileStatusChange_actionProfileUnavailableForPrivateProfile() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         UserId managedUser = UserId.of(mManagedUser);
@@ -603,9 +625,6 @@
                 mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
 
         List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser);
-        Map<UserId, Boolean> expectedCanForwardToProfileIdMapAfterIntent = new HashMap<>();
-        expectedCanForwardToProfileIdMapAfterIntent.put(currentUser, true);
-        expectedCanForwardToProfileIdMapAfterIntent.put(managedUser, true);
 
         String action = Intent.ACTION_PROFILE_UNAVAILABLE;
         mUserManagerState.onProfileActionStatusChange(action, privateUser);
@@ -615,19 +634,14 @@
                 .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
         assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
                 .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
-        assertWithMessage(
-                "CanForwardToLabelMap should not be same before and after receiving intent: "
-                        + action)
-                .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isNotEqualTo(
-                        canForwardToProfileIdMapBeforeIntent);
-        assertWithMessage(
-                "Unexpected changes to canForwardToProfileIdMap on receiving intent: " + action)
+        assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
+                + action)
                 .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
-                        expectedCanForwardToProfileIdMapAfterIntent);
+                        canForwardToProfileIdMapBeforeIntent);
     }
 
     @Test
-    public void testOnProfileStatusChange_actionProfileAvailableOnPrivateProfile() {
+    public void testOnProfileStatusChange_actionProfileAvailable_profileInitialised() {
         if (!SdkLevel.isAtLeastV()) return;
         UserId currentUser = UserId.of(mSystemUser);
         UserId managedUser = UserId.of(mManagedUser);
@@ -637,9 +651,54 @@
                 PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
                 mMockResolveInfoList);
         initializeUserManagerState(currentUser,
-                Lists.newArrayList(mSystemUser, mManagedUser));
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
 
-        // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will empty
+        // initialising the userIds list and canForwardToProfileIdMap
+        mUserManagerState.getUserIds();
+        mUserManagerState.getCanForwardToProfileIdMap(mMockIntent);
+
+        // Making the private profile unavailable after it has been initialised
+        mUserManagerState.onProfileActionStatusChange(Intent.ACTION_PROFILE_UNAVAILABLE,
+                privateUser);
+
+        List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
+        Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
+                mUserManagerState.getCanForwardToProfileIdMap(mMockIntent));
+
+        List<UserId> expectedUserIdsAfterIntent = Lists.newArrayList(currentUser, managedUser,
+                privateUser);
+
+        String action = Intent.ACTION_PROFILE_AVAILABLE;
+        mUserManagerState.onProfileActionStatusChange(action, privateUser);
+
+        assertWithMessage(
+                "UserIds list should not be same before and after receiving intent: " + action)
+                .that(mUserManagerState.getUserIds()).isNotEqualTo(userIdsBeforeIntent);
+        assertWithMessage("Unexpected changes to user id list on receiving intent: " + action)
+                .that(mUserManagerState.getUserIds()).isEqualTo(expectedUserIdsAfterIntent);
+        assertWithMessage("CanForwardToLabelMap should be same before and after receiving intent: "
+                + action)
+                .that(mUserManagerState.getCanForwardToProfileIdMap(mMockIntent)).isEqualTo(
+                        canForwardToProfileIdMapBeforeIntent);
+    }
+
+    @Test
+    public void testOnProfileStatusChange_actionProfileAvailable_profileNotInitialised() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        UserId managedUser = UserId.of(mManagedUser);
+        UserId privateUser = UserId.of(mPrivateUser);
+        final List<ResolveInfo> mMockResolveInfoList = Lists.newArrayList(mMockInfo1, mMockInfo2);
+        when(mMockPackageManager.queryIntentActivitiesAsUser(mMockIntent,
+                PackageManager.MATCH_DEFAULT_ONLY, mSystemUser)).thenReturn(
+                mMockResolveInfoList);
+
+        // Private user will not be initialised if it is in quiet mode
+        when(mMockUserManager.isQuietModeEnabled(mPrivateUser)).thenReturn(true);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+
+        // UserManagerState#mUserId and UserManagerState#mCanForwardToProfileIdMap will be empty
         // by default if the getters of these member variables have not been called
         List<UserId> userIdsBeforeIntent = new ArrayList<>(mUserManagerState.getUserIds());
         Map<UserId, Boolean> canForwardToProfileIdMapBeforeIntent = new HashMap<>(
@@ -671,6 +730,103 @@
                         expectedCanForwardToProfileIdMapAfterIntent);
     }
 
+    @Test
+    public void testGetUserIdToLabelMap_systemUserAndManagedUser_PreV() {
+        if (SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser));
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
+                    PERSONAL);
+            when(devicePolicyResourcesManager.getString(eq(WORK_TAB), any())).thenReturn(WORK);
+        }
+
+        Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
+
+        assertWithMessage("Incorrect label returned for user id " + mSystemUser)
+                .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
+        assertWithMessage("Incorrect label returned for user id " + mManagedUser)
+                .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
+    }
+
+    @Test
+    public void testGetUserIdToLabelMap_systemUserManagedUserPrivateUser_PostV() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getString(eq(PERSONAL_TAB), any())).thenReturn(
+                    PERSONAL);
+        }
+        UserManager managedUserManager = getUserManagerForManagedUser();
+        UserManager privateUserManager = getUserManagerForPrivateUser();
+        when(managedUserManager.getProfileLabel()).thenReturn(WORK);
+        when(privateUserManager.getProfileLabel()).thenReturn(PRIVATE);
+
+        Map<UserId, String> userIdToLabelMap = mUserManagerState.getUserIdToLabelMap();
+
+        assertWithMessage("Incorrect label returned for user id " + mSystemUser)
+                .that(userIdToLabelMap.get(UserId.of(mSystemUser))).isEqualTo(PERSONAL);
+        assertWithMessage("Incorrect label returned for user id " + mManagedUser)
+                .that(userIdToLabelMap.get(UserId.of(mManagedUser))).isEqualTo(WORK);
+        assertWithMessage("Incorrect label returned for user id " + mPrivateUser)
+                .that(userIdToLabelMap.get(UserId.of(mPrivateUser))).isEqualTo(PRIVATE);
+    }
+
+    @Test
+    public void testGetUserIdToBadgeMap_systemUserManagedUser_PreV() {
+        if (SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser));
+        Drawable workBadge = mMockContext.getDrawable(R.drawable.ic_briefcase);
+        if (SdkLevel.isAtLeastT()) {
+            DevicePolicyResourcesManager devicePolicyResourcesManager = mock(
+                    DevicePolicyResourcesManager.class);
+            when(mDevicePolicyManager.getResources()).thenReturn(devicePolicyResourcesManager);
+            when(devicePolicyResourcesManager.getDrawable(eq(WORK_PROFILE_ICON), eq(SOLID_COLORED),
+                    any())).thenReturn(workBadge);
+        }
+
+        Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
+
+        assertWithMessage("There should be no badge present for personal user")
+                .that(userIdToBadgeMap.containsKey(UserId.of(mSystemUser))).isFalse();
+        assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
+                .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
+    }
+
+    @Test
+    public void testGetUserIdToBadgeMap_systemUserManagedUserPrivateUser_PostV() {
+        if (!SdkLevel.isAtLeastV()) return;
+        UserId currentUser = UserId.of(mSystemUser);
+        initializeUserManagerState(currentUser,
+                Lists.newArrayList(mSystemUser, mManagedUser, mPrivateUser));
+        Drawable workBadge = mock(Drawable.class);
+        Drawable privateBadge = mock(Drawable.class);
+        UserManager managedUserManager = getUserManagerForManagedUser();
+        UserManager privateUserManager = getUserManagerForPrivateUser();
+        when(managedUserManager.getUserBadge()).thenReturn(workBadge);
+        when(privateUserManager.getUserBadge()).thenReturn(privateBadge);
+
+        Map<UserId, Drawable> userIdToBadgeMap = mUserManagerState.getUserIdToBadgeMap();
+
+        assertWithMessage("There should be no badge present for personal user")
+                .that(userIdToBadgeMap.get(UserId.of(mSystemUser))).isNull();
+        assertWithMessage("Incorrect badge returned for user id " + mManagedUser)
+                .that(userIdToBadgeMap.get(UserId.of(mManagedUser))).isEqualTo(workBadge);
+        assertWithMessage("Incorrect badge returned for user id " + mPrivateUser)
+                .that(userIdToBadgeMap.get(UserId.of(mPrivateUser))).isEqualTo(privateBadge);
+    }
+
     private void initializeUserManagerState(UserId current, List<UserHandle> usersOnDevice) {
         when(mMockUserManager.getUserProfiles()).thenReturn(usersOnDevice);
         TestConfigStore testConfigStore = new TestConfigStore();
@@ -678,4 +834,24 @@
         mUserManagerState = new UserManagerState.RuntimeUserManagerState(mMockContext, current,
                 true, testConfigStore);
     }
+
+    private UserManager getUserManagerForManagedUser() {
+        Context managedUserContext = mock(Context.class);
+        when(mMockContext.createContextAsUser(mManagedUser, 0)).thenReturn(managedUserContext);
+        UserManager managedUserManager = mock(UserManager.class);
+        when(managedUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn("managedUserManager");
+        when(managedUserContext.getSystemService(UserManager.class)).thenReturn(managedUserManager);
+        return managedUserManager;
+    }
+
+    private UserManager getUserManagerForPrivateUser() {
+        Context privateUserContext = mock(Context.class);
+        when(mMockContext.createContextAsUser(mPrivateUser, 0)).thenReturn(privateUserContext);
+        UserManager privateUserManager = mock(UserManager.class);
+        when(privateUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn("privateUserManager");
+        when(privateUserContext.getSystemService(UserManager.class)).thenReturn(privateUserManager);
+        return privateUserManager;
+    }
 }