Move the selection menu from the ActionBar to NavigationView

In the >600dp layout, the MaterialToolbar has been moved to be
adjacent to the sidebar. Previously the ActionBar was used to replace
the existing toolbar, unfortunately on the new layout this doesn't
work as it spans the full width.

To get around this (and maintain compatibility when the flag is on or
off) we instead inflate the activity menu OR the action mode menu
depending on whether there is a selection. This is a temporary measure
to not have to refactor the click handlers too much. In the final
product the ActionModeController will be removed and the click
handlers consolidated into a single place.

Bug: 383669583
Test: atest DocumentsUIGoogleTests
Flag: com.android.documentsui.flags.use_material3

Change-Id: Ida4675dafbc2096c986067ba4fbe5a695a8e243b

13 files changed, 394 insertions, 127 deletions

diff --git a/res/flag(com.android.documentsui.flags.use_material3)/drawable/ic_cancel.xml b/res/flag(com.android.documentsui.flags.use_material3)/drawable/ic_cancel.xml
new file mode 100644
index 000000000..f47cd4f6b
--- /dev/null
+++ b/res/flag(com.android.documentsui.flags.use_material3)/drawable/ic_cancel.xml
@@ -0,0 +1,24 @@
+<!--
+    Copyright (C) 2025 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportHeight="960"
+    android:viewportWidth="960">
+    <path
+        android:fillColor="?attr/colorOnSurface"
+        android:pathData="M256,760L200,704L424,480L200,256L256,200L480,424L704,200L760,256L536,480L760,704L704,760L480,536L256,760Z" />
+</vector>
diff --git a/res/flag(com.android.documentsui.flags.use_material3)/menu/action_mode_menu.xml b/res/flag(com.android.documentsui.flags.use_material3)/menu/action_mode_menu.xml
index db4d581c1..c9c139cc9 100644
--- a/res/flag(com.android.documentsui.flags.use_material3)/menu/action_mode_menu.xml
+++ b/res/flag(com.android.documentsui.flags.use_material3)/menu/action_mode_menu.xml
@@ -14,72 +14,88 @@
      limitations under the License.
 -->
 
-<menu xmlns:android="http://schemas.android.com/apk/res/android">
+<menu
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto">
     <item
         android:id="@+id/action_menu_open_with"
         android:title="@string/menu_open_with"
-        android:showAsAction="never" />
+        android:showAsAction="never"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_share"
         android:icon="@drawable/ic_menu_share"
         android:title="@string/menu_share"
-        android:showAsAction="always" />
+        android:showAsAction="always"
+        app:showAsAction="always" />
     <item
         android:id="@+id/action_menu_delete"
         android:icon="@drawable/ic_menu_delete"
         android:title="@string/menu_delete"
-        android:showAsAction="always" />
+        android:showAsAction="always"
+        app:showAsAction="always" />
     <item
         android:id="@+id/action_menu_sort"
         android:icon="@drawable/ic_sort"
         android:title="@string/menu_sort"
-        android:showAsAction="never" />
+        android:showAsAction="never"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_select"
         android:title="@string/menu_select"
-        android:showAsAction="always" />
+        android:showAsAction="always"
+        app:showAsAction="always" />
     <item
         android:id="@+id/action_menu_select_all"
         android:title="@string/menu_select_all"
-        android:showAsAction="never" />
+        android:showAsAction="never"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_deselect_all"
         android:title="@string/menu_deselect_all"
-        android:showAsAction="never" />
+        android:showAsAction="never"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_copy_to"
         android:title="@string/menu_copy"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_extract_to"
         android:title="@string/menu_extract"
         android:icon="@drawable/ic_menu_extract"
         android:showAsAction="always"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="always" />
     <item
         android:id="@+id/action_menu_move_to"
         android:title="@string/menu_move"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_compress"
         android:title="@string/menu_compress"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_rename"
         android:title="@string/menu_rename"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_inspect"
         android:title="@string/menu_inspect"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
     <item
         android:id="@+id/action_menu_view_in_owner"
         android:title="@string/menu_view_in_owner"
         android:showAsAction="never"
-        android:visible="false" />
+        android:visible="false"
+        app:showAsAction="never" />
 </menu>
diff --git a/src/com/android/documentsui/ActionModeController.java b/src/com/android/documentsui/ActionModeController.java
index 39a50abf9..6259a8bba 100644
--- a/src/com/android/documentsui/ActionModeController.java
+++ b/src/com/android/documentsui/ActionModeController.java
@@ -39,6 +39,8 @@ import com.android.documentsui.ui.MessageBuilder;
 
 /**
  * A controller that listens to selection changes and manages life cycles of action modes.
+ * TODO(b/379776735): This class (and action mode in general) is no longer in use when the
+ * use_material3 flag is enabled. Remove the class once the flag is rolled out.
  */
 public class ActionModeController extends SelectionObserver<String>
         implements ActionMode.Callback, ActionModeAddons {
diff --git a/src/com/android/documentsui/BaseActivity.java b/src/com/android/documentsui/BaseActivity.java
index 41fc5182f..9421d2325 100644
--- a/src/com/android/documentsui/BaseActivity.java
+++ b/src/com/android/documentsui/BaseActivity.java
@@ -401,13 +401,27 @@ public abstract class BaseActivity
     private NavigationViewManager getNavigationViewManager(Breadcrumb breadcrumb,
             View profileTabsContainer) {
         if (mConfigStore.isPrivateSpaceInDocsUIEnabled()) {
-            return new NavigationViewManager(this, mDrawer, mState, this, breadcrumb,
-                    profileTabsContainer, DocumentsApplication.getUserManagerState(this),
-                    mConfigStore);
+            return new NavigationViewManager(
+                    this,
+                    mDrawer,
+                    mState,
+                    this,
+                    breadcrumb,
+                    profileTabsContainer,
+                    DocumentsApplication.getUserManagerState(this),
+                    mConfigStore,
+                    mInjector);
         }
-        return new NavigationViewManager(this, mDrawer, mState, this, breadcrumb,
-                profileTabsContainer, DocumentsApplication.getUserIdManager(this),
-                mConfigStore);
+        return new NavigationViewManager(
+                this,
+                mDrawer,
+                mState,
+                this,
+                breadcrumb,
+                profileTabsContainer,
+                DocumentsApplication.getUserIdManager(this),
+                mConfigStore,
+                mInjector);
     }
 
     public void onPreferenceChanged(String pref) {
@@ -421,14 +435,21 @@ public abstract class BaseActivity
     protected void onPostCreate(Bundle savedInstanceState) {
         super.onPostCreate(savedInstanceState);
 
-        mRootsMonitor = new RootsMonitor<>(
-                this,
-                mInjector.actions,
-                mProviders,
-                mDocs,
-                mState,
-                mSearchManager,
-                mInjector.actionModeController::finishActionMode);
+        Runnable finishActionMode =
+                (useMaterial3())
+                        ? mNavigator::closeSelectionBar
+                        : mInjector.actionModeController::finishActionMode;
+
+        mRootsMonitor =
+                new RootsMonitor<>(
+                        this,
+                        mInjector.actions,
+                        mProviders,
+                        mDocs,
+                        mState,
+                        mSearchManager,
+                        finishActionMode);
+
         mRootsMonitor.start();
     }
 
@@ -440,6 +461,13 @@ public abstract class BaseActivity
 
     @Override
     public boolean onCreateOptionsMenu(Menu menu) {
+        if (useMaterial3()) {
+            // In Material3 the menu is now inflated in the `NavigationViewMenu`. This is currently
+            // to allow for us to inflate between the action_menu and the activity menu. Once the
+            // Material 3 flag is removed, the menus will be merged and we can rely on this single
+            // inflation point.
+            return super.onCreateOptionsMenu(menu);
+        }
         boolean showMenu = super.onCreateOptionsMenu(menu);
 
         getMenuInflater().inflate(R.menu.activity, menu);
@@ -463,11 +491,13 @@ public abstract class BaseActivity
     @CallSuper
     public boolean onPrepareOptionsMenu(Menu menu) {
         super.onPrepareOptionsMenu(menu);
-        mSearchManager.showMenu(mState.stack);
         // Remove the subMenu when material3 is launched b/379776735.
         if (useMaterial3()) {
-            mInjector.menuManager.updateSubMenu(null);
+            if (mNavigator != null) {
+                mNavigator.updateActionMenu();
+            }
         } else {
+            mSearchManager.showMenu(mState.stack);
             final ActionMenuView subMenuView = findViewById(R.id.sub_menu);
             mInjector.menuManager.updateSubMenu(subMenuView.getMenu());
         }
@@ -569,7 +599,11 @@ public abstract class BaseActivity
             return;
         }
 
-        mInjector.actionModeController.finishActionMode();
+        if (useMaterial3()) {
+            mNavigator.closeSelectionBar();
+        } else {
+            mInjector.actionModeController.finishActionMode();
+        }
         mSortController.onViewModeChanged(mState.derivedMode);
 
         // Set summary header's visibility. Only recents and downloads root may have summary in
@@ -668,6 +702,10 @@ public abstract class BaseActivity
         mNavigator.update();
     }
 
+    public final NavigationViewManager getNavigator() {
+        return mNavigator;
+    }
+
     @Override
     public void restoreRootAndDirectory() {
         // We're trying to restore stuff in document stack from saved instance. If we didn't have a
diff --git a/src/com/android/documentsui/Injector.java b/src/com/android/documentsui/Injector.java
index 82cbfcccb..5fd716a97 100644
--- a/src/com/android/documentsui/Injector.java
+++ b/src/com/android/documentsui/Injector.java
@@ -15,6 +15,8 @@
  */
 package com.android.documentsui;
 
+import static com.android.documentsui.flags.Flags.useMaterial3;
+
 import static java.lang.annotation.ElementType.FIELD;
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
@@ -127,6 +129,9 @@ public class Injector<T extends ActionHandler> {
 
     public final ActionModeController getActionModeController(
             SelectionDetails selectionDetails, EventHandler<MenuItem> menuItemClicker) {
+        if (useMaterial3()) {
+            return null;
+        }
         return actionModeController.reset(selectionDetails, menuItemClicker);
     }
 
diff --git a/src/com/android/documentsui/NavigationViewManager.java b/src/com/android/documentsui/NavigationViewManager.java
index c376c86db..c6739baaf 100644
--- a/src/com/android/documentsui/NavigationViewManager.java
+++ b/src/com/android/documentsui/NavigationViewManager.java
@@ -17,6 +17,7 @@
 package com.android.documentsui;
 
 import static com.android.documentsui.base.SharedMinimal.VERBOSE;
+import static com.android.documentsui.flags.Flags.useMaterial3;
 
 import android.content.res.Resources;
 import android.content.res.TypedArray;
@@ -24,6 +25,7 @@ import android.graphics.Outline;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.util.Log;
+import android.view.MenuItem;
 import android.view.View;
 import android.view.ViewOutlineProvider;
 import android.view.Window;
@@ -34,7 +36,10 @@ import androidx.annotation.ColorRes;
 import androidx.annotation.Nullable;
 import androidx.appcompat.widget.Toolbar;
 import androidx.core.content.ContextCompat;
+import androidx.recyclerview.selection.SelectionTracker;
 
+import com.android.documentsui.Injector.Injected;
+import com.android.documentsui.base.EventHandler;
 import com.android.documentsui.base.RootInfo;
 import com.android.documentsui.base.State;
 import com.android.documentsui.base.UserId;
@@ -47,10 +52,9 @@ import com.google.android.material.appbar.CollapsingToolbarLayout;
 
 import java.util.function.IntConsumer;
 
-/**
- * A facade over the portions of the app and drawer toolbars.
- */
-public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListener {
+/** A facade over the portions of the app and drawer toolbars. */
+public class NavigationViewManager extends SelectionTracker.SelectionObserver<String>
+        implements AppBarLayout.OnOffsetChangedListener {
 
     private static final String TAG = "NavigationViewManager";
 
@@ -69,10 +73,11 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
     private final ViewOutlineProvider mSearchBarOutlineProvider;
     private final boolean mShowSearchBar;
     private final ConfigStore mConfigStore;
-
+    @Injected private final Injector<?> mInjector;
     private boolean mIsActionModeActivated = false;
-    @ColorRes
-    private int mDefaultStatusBarColorResId;
+    @ColorRes private int mDefaultStatusBarColorResId;
+    private MenuManager.SelectionDetails mSelectionDetails;
+    private EventHandler<MenuItem> mActionMenuItemClicker;
 
     public NavigationViewManager(
             BaseActivity activity,
@@ -82,9 +87,19 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
             Breadcrumb breadcrumb,
             View tabLayoutContainer,
             UserIdManager userIdManager,
-            ConfigStore configStore) {
-        this(activity, drawer, state, env, breadcrumb, tabLayoutContainer, userIdManager, null,
-                configStore);
+            ConfigStore configStore,
+            Injector injector) {
+        this(
+                activity,
+                drawer,
+                state,
+                env,
+                breadcrumb,
+                tabLayoutContainer,
+                userIdManager,
+                null,
+                configStore,
+                injector);
     }
 
     public NavigationViewManager(
@@ -95,9 +110,19 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
             Breadcrumb breadcrumb,
             View tabLayoutContainer,
             UserManagerState userManagerState,
-            ConfigStore configStore) {
-        this(activity, drawer, state, env, breadcrumb, tabLayoutContainer, null, userManagerState,
-                configStore);
+            ConfigStore configStore,
+            Injector injector) {
+        this(
+                activity,
+                drawer,
+                state,
+                env,
+                breadcrumb,
+                tabLayoutContainer,
+                null,
+                userManagerState,
+                configStore,
+                injector);
     }
 
     public NavigationViewManager(
@@ -109,7 +134,8 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
             View tabLayoutContainer,
             UserIdManager userIdManager,
             UserManagerState userManagerState,
-            ConfigStore configStore) {
+            ConfigStore configStore,
+            Injector injector) {
 
         mActivity = activity;
         mToolbar = activity.findViewById(R.id.toolbar);
@@ -120,6 +146,7 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
         mBreadcrumb = breadcrumb;
         mBreadcrumb.setup(env, state, this::onNavigationItemSelected);
         mConfigStore = configStore;
+        mInjector = injector;
         mProfileTabs =
                 getProfileTabs(tabLayoutContainer, userIdManager, userManagerState, activity);
 
@@ -130,6 +157,15 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
                         onNavigationIconClicked();
                     }
                 });
+        if (useMaterial3()) {
+            mToolbar.setOnMenuItemClickListener(
+                    new Toolbar.OnMenuItemClickListener() {
+                        @Override
+                        public boolean onMenuItemClick(MenuItem menuItem) {
+                            return onToolbarMenuItemClicked(menuItem);
+                        }
+                    });
+        }
         mSearchBarView = activity.findViewById(R.id.searchbar_title);
         mCollapsingBarLayout = activity.findViewById(R.id.collapsing_toolbar);
         mDefaultActionBarBackground = mToolbar.getBackground();
@@ -225,11 +261,21 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
     }
 
     private void onNavigationIconClicked() {
-        if (mDrawer.isPresent()) {
+        if (useMaterial3() && inSelectionMode()) {
+            closeSelectionBar();
+        } else if (mDrawer.isPresent()) {
             mDrawer.setOpen(true);
         }
     }
 
+    private boolean onToolbarMenuItemClicked(MenuItem menuItem) {
+        if (inSelectionMode()) {
+            mActionMenuItemClicker.accept(menuItem);
+            return true;
+        }
+        return mActivity.onOptionsItemSelected(menuItem);
+    }
+
     void onNavigationItemSelected(int position) {
         boolean changed = false;
         while (mState.stack.size() > position + 1) {
@@ -271,15 +317,86 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
             mBreadcrumb.show(false);
             mToolbar.setTitle(null);
             mSearchBarView.setVisibility(View.VISIBLE);
-        } else {
-            mSearchBarView.setVisibility(View.GONE);
-            String title = mState.stack.size() <= 1
-                    ? mEnv.getCurrentRoot().title : mState.stack.getTitle();
-            if (VERBOSE) Log.v(TAG, "New toolbar title is: " + title);
-            mToolbar.setTitle(title);
-            mBreadcrumb.show(true);
-            mBreadcrumb.postUpdate();
+            return;
+        }
+
+        mSearchBarView.setVisibility(View.GONE);
+
+        if (useMaterial3()) {
+            updateActionMenu();
+            if (inSelectionMode()) {
+                final int quantity = mInjector.selectionMgr.getSelection().size();
+                final String title =
+                        mToolbar.getContext()
+                                .getResources()
+                                .getQuantityString(R.plurals.elements_selected, quantity, quantity);
+                mToolbar.setTitle(title);
+                mActivity.getWindow().setTitle(title);
+                mToolbar.setNavigationIcon(R.drawable.ic_cancel);
+                mToolbar.setNavigationContentDescription(android.R.string.cancel);
+                return;
+            }
+        }
+
+        String title =
+                mState.stack.size() <= 1 ? mEnv.getCurrentRoot().title : mState.stack.getTitle();
+        if (VERBOSE) Log.v(TAG, "New toolbar title is: " + title);
+        mToolbar.setTitle(title);
+        mBreadcrumb.show(true);
+        mBreadcrumb.postUpdate();
+    }
+
+    @Override
+    public void onSelectionChanged() {
+        update();
+    }
+
+    /** Identifies if the `NavigationViewManager` is in selection mode or not. */
+    public boolean inSelectionMode() {
+        return mInjector != null
+                && mInjector.selectionMgr != null
+                && mInjector.selectionMgr.hasSelection();
+    }
+
+    private boolean hasActionMenu() {
+        return mToolbar.getMenu().findItem(R.id.action_menu_open_with) != null;
+    }
+
+    /** Updates the action menu based on whether a selection is currently being made or not. */
+    public void updateActionMenu() {
+        // For the first start up of the application, the menu might not exist at all but we also
+        // don't want to inflate the menu multiple times. So along with checking if the expected
+        // menu is already inflated, validate that a menu exists at all as well.
+        boolean isMenuInflated = mToolbar.getMenu() != null && mToolbar.getMenu().size() > 0;
+        if (inSelectionMode()) {
+            if (!isMenuInflated || !hasActionMenu()) {
+                mToolbar.getMenu().clear();
+                mToolbar.inflateMenu(R.menu.action_mode_menu);
+                mToolbar.invalidateMenu();
+            }
+            mInjector.menuManager.updateActionMenu(mToolbar.getMenu(), mSelectionDetails);
+            return;
+        }
+
+        if (!isMenuInflated || hasActionMenu()) {
+            mToolbar.getMenu().clear();
+            mToolbar.inflateMenu(R.menu.activity);
+            mToolbar.invalidateMenu();
+            boolean fullBarSearch =
+                    mActivity.getResources().getBoolean(R.bool.full_bar_search_view);
+            boolean showSearchBar = mActivity.getResources().getBoolean(R.bool.show_search_bar);
+            mInjector.searchManager.install(mToolbar.getMenu(), fullBarSearch, showSearchBar);
         }
+        mInjector.menuManager.updateOptionMenu(mToolbar.getMenu());
+        mInjector.searchManager.showMenu(mState.stack);
+    }
+
+    /** Everytime a selection is made, update the selection. */
+    public void updateSelection(
+            MenuManager.SelectionDetails selectionDetails,
+            EventHandler<MenuItem> actionMenuItemClicker) {
+        mSelectionDetails = selectionDetails;
+        mActionMenuItemClicker = actionMenuItemClicker;
     }
 
     private void updateScrollFlag() {
@@ -361,6 +478,11 @@ public class NavigationViewManager implements AppBarLayout.OnOffsetChangedListen
         mDrawer.setOpen(open);
     }
 
+    /** Helper method to close the selection bar. */
+    public void closeSelectionBar() {
+        mInjector.selectionMgr.clearSelection();
+    }
+
     interface Breadcrumb {
         void setup(Environment env, State state, IntConsumer listener);
 
diff --git a/src/com/android/documentsui/dirlist/DirectoryFragment.java b/src/com/android/documentsui/dirlist/DirectoryFragment.java
index ffa912c51..078498153 100644
--- a/src/com/android/documentsui/dirlist/DirectoryFragment.java
+++ b/src/com/android/documentsui/dirlist/DirectoryFragment.java
@@ -22,6 +22,7 @@ import static com.android.documentsui.base.SharedMinimal.VERBOSE;
 import static com.android.documentsui.base.State.MODE_GRID;
 import static com.android.documentsui.base.State.MODE_LIST;
 import static com.android.documentsui.flags.Flags.desktopFileHandling;
+import static com.android.documentsui.flags.Flags.useMaterial3;
 
 import android.app.ActivityManager;
 import android.content.BroadcastReceiver;
@@ -610,11 +611,16 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
         new RefreshHelper(mRefreshLayout::setEnabled)
                 .attach(mRecView);
 
-        mActionModeController = mInjector.getActionModeController(
-                mSelectionMetadata,
-                this::handleMenuItemClick);
-
-        mSelectionMgr.addObserver(mActionModeController);
+        if (useMaterial3()) {
+            mSelectionMgr.addObserver(mActivity.getNavigator());
+            mActivity.getNavigator().updateSelection(mSelectionMetadata, this::handleMenuItemClick);
+        } else {
+            mActionModeController =
+                    mInjector.getActionModeController(
+                            mSelectionMetadata, this::handleMenuItemClick);
+            assert (mActionModeController != null);
+            mSelectionMgr.addObserver(mActionModeController);
+        }
 
         mProfileTabsController = mInjector.profileTabsController;
         mSelectionMgr.addObserver(mProfileTabsController);
@@ -917,6 +923,14 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
         }
     }
 
+    private void closeSelectionBar() {
+        if (useMaterial3()) {
+            mActivity.getNavigator().closeSelectionBar();
+        } else {
+            mActionModeController.finishActionMode();
+        }
+    }
+
     private boolean handleMenuItemClick(MenuItem item) {
         if (mInjector.pickResult != null) {
             mInjector.pickResult.increaseActionCount();
@@ -937,7 +951,7 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
             // hidden unless the desktopFileHandling flag is enabled, in which case the menu item
             // will be handled by the condition above.
             openDocuments(selection);
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             return true;
         } else if (id == R.id.action_menu_open_with || id == R.id.dir_menu_open_with) {
             showChooserForDoc(selection);
@@ -957,14 +971,14 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
             transferDocuments(selection, null, FileOperationService.OPERATION_COPY);
             // TODO: Only finish selection mode if copy-to is not canceled.
             // Need to plum down into handling the way we do with deleteDocuments.
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             return true;
         } else if (id == R.id.action_menu_compress || id == R.id.dir_menu_compress) {
             transferDocuments(selection, mState.stack,
                     FileOperationService.OPERATION_COMPRESS);
             // TODO: Only finish selection mode if compress is not canceled.
             // Need to plum down into handling the way we do with deleteDocuments.
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             return true;
 
             // TODO: Implement extract (to the current directory).
@@ -972,7 +986,7 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
             transferDocuments(selection, null, FileOperationService.OPERATION_EXTRACT);
             // TODO: Only finish selection mode if compress-to is not canceled.
             // Need to plum down into handling the way we do with deleteDocuments.
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             return true;
         } else if (id == R.id.action_menu_move_to) {
             if (mModel.hasDocuments(selection, DocumentFilters.NOT_MOVABLE)) {
@@ -980,11 +994,11 @@ public class DirectoryFragment extends Fragment implements SwipeRefreshLayout.On
                 return true;
             }
             // Exit selection mode first, so we avoid deselecting deleted documents.
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             transferDocuments(selection, null, FileOperationService.OPERATION_MOVE);
             return true;
         } else if (id == R.id.action_menu_inspect || id == R.id.dir_menu_inspect) {
-            mActionModeController.finishActionMode();
+            closeSelectionBar();
             assert selection.size() <= 1;
             DocumentInfo doc = selection.isEmpty()
                     ? mActivity.getCurrentDirectory()
diff --git a/src/com/android/documentsui/files/ActionHandler.java b/src/com/android/documentsui/files/ActionHandler.java
index af5ef5bbc..33057d140 100644
--- a/src/com/android/documentsui/files/ActionHandler.java
+++ b/src/com/android/documentsui/files/ActionHandler.java
@@ -19,7 +19,6 @@ package com.android.documentsui.files;
 import static android.content.ContentResolver.wrap;
 
 import static com.android.documentsui.base.SharedMinimal.DEBUG;
-import com.android.documentsui.flags.Flags;
 
 import android.app.DownloadManager;
 import android.content.ActivityNotFoundException;
@@ -68,6 +67,7 @@ import com.android.documentsui.clipping.ClipStore;
 import com.android.documentsui.clipping.DocumentClipper;
 import com.android.documentsui.clipping.UrisSupplier;
 import com.android.documentsui.dirlist.AnimationView;
+import com.android.documentsui.flags.Flags;
 import com.android.documentsui.inspector.InspectorActivity;
 import com.android.documentsui.queries.SearchViewManager;
 import com.android.documentsui.roots.ProvidersAccess;
@@ -98,6 +98,7 @@ public class ActionHandler<T extends FragmentActivity & AbstractActionHandler.Co
     private final DocumentClipper mClipper;
     private final ClipStore mClipStore;
     private final DragAndDropManager mDragAndDropManager;
+    private final Runnable mCloseSelectionBar;
 
     ActionHandler(
             T activity,
@@ -106,7 +107,8 @@ public class ActionHandler<T extends FragmentActivity & AbstractActionHandler.Co
             DocumentsAccess docs,
             SearchViewManager searchMgr,
             Lookup<String, Executor> executors,
-            ActionModeAddons actionModeAddons,
+            @Nullable ActionModeAddons actionModeAddons,
+            Runnable closeSelectionBar,
             DocumentClipper clipper,
             ClipStore clipStore,
             DragAndDropManager dragAndDropManager,
@@ -115,6 +117,7 @@ public class ActionHandler<T extends FragmentActivity & AbstractActionHandler.Co
         super(activity, state, providers, docs, searchMgr, executors, injector);
 
         mActionModeAddons = actionModeAddons;
+        mCloseSelectionBar = closeSelectionBar;
         mFeatures = injector.features;
         mConfig = injector.config;
         mClipper = clipper;
@@ -230,8 +233,12 @@ public class ActionHandler<T extends FragmentActivity & AbstractActionHandler.Co
 
     @Override
     public void springOpenDirectory(DocumentInfo doc) {
-        assert(doc.isDirectory());
-        mActionModeAddons.finishActionMode();
+        assert (doc.isDirectory());
+        if (Flags.useMaterial3()) {
+            mCloseSelectionBar.run();
+        } else {
+            mActionModeAddons.finishActionMode();
+        }
         openContainerDocument(doc);
     }
 
@@ -323,7 +330,11 @@ public class ActionHandler<T extends FragmentActivity & AbstractActionHandler.Co
             return;
         }
 
-        mActionModeAddons.finishActionMode();
+        if (Flags.useMaterial3()) {
+            mCloseSelectionBar.run();
+        } else {
+            mActionModeAddons.finishActionMode();
+        }
 
         List<Uri> uris = new ArrayList<>(docs.size());
         for (DocumentInfo doc : docs) {
diff --git a/src/com/android/documentsui/files/FilesActivity.java b/src/com/android/documentsui/files/FilesActivity.java
index 4bb7c1ac5..ff9d30106 100644
--- a/src/com/android/documentsui/files/FilesActivity.java
+++ b/src/com/android/documentsui/files/FilesActivity.java
@@ -140,25 +140,30 @@ public class FilesActivity extends BaseActivity implements AbstractActionHandler
                 mInjector.getModel()::getItemUri,
                 mInjector.getModel()::getItemCount);
 
-        mInjector.actionModeController = new ActionModeController(
-                this,
-                mInjector.selectionMgr,
-                mNavigator,
-                mInjector.menuManager,
-                mInjector.messages);
+        if (!useMaterial3()) {
+            mInjector.actionModeController =
+                    new ActionModeController(
+                            this,
+                            mInjector.selectionMgr,
+                            mNavigator,
+                            mInjector.menuManager,
+                            mInjector.messages);
+        }
 
-        mInjector.actions = new ActionHandler<>(
-                this,
-                mState,
-                mProviders,
-                mDocs,
-                mSearchManager,
-                ProviderExecutor::forAuthority,
-                mInjector.actionModeController,
-                clipper,
-                DocumentsApplication.getClipStore(this),
-                DocumentsApplication.getDragAndDropManager(this),
-                mInjector);
+        mInjector.actions =
+                new ActionHandler<>(
+                        this,
+                        mState,
+                        mProviders,
+                        mDocs,
+                        mSearchManager,
+                        ProviderExecutor::forAuthority,
+                        mInjector.actionModeController,
+                        getNavigator()::closeSelectionBar,
+                        clipper,
+                        DocumentsApplication.getClipStore(this),
+                        DocumentsApplication.getDragAndDropManager(this),
+                        mInjector);
 
         mInjector.searchManager = mSearchManager;
 
@@ -327,7 +332,9 @@ public class FilesActivity extends BaseActivity implements AbstractActionHandler
     @Override
     public boolean onPrepareOptionsMenu(Menu menu) {
         super.onPrepareOptionsMenu(menu);
-        mInjector.menuManager.updateOptionMenu(menu);
+        if (!useMaterial3()) {
+            mInjector.menuManager.updateOptionMenu(menu);
+        }
         return true;
     }
 
diff --git a/src/com/android/documentsui/picker/PickActivity.java b/src/com/android/documentsui/picker/PickActivity.java
index 481b67e77..51b35e155 100644
--- a/src/com/android/documentsui/picker/PickActivity.java
+++ b/src/com/android/documentsui/picker/PickActivity.java
@@ -128,12 +128,15 @@ public class PickActivity extends BaseActivity implements ActionHandler.Addons {
                 new DirectoryDetails(this),
                 mInjector.getModel()::getItemCount);
 
-        mInjector.actionModeController = new ActionModeController(
-                this,
-                mInjector.selectionMgr,
-                mNavigator,
-                mInjector.menuManager,
-                mInjector.messages);
+        if (!useMaterial3()) {
+            mInjector.actionModeController =
+                    new ActionModeController(
+                            this,
+                            mInjector.selectionMgr,
+                            mNavigator,
+                            mInjector.menuManager,
+                            mInjector.messages);
+        }
 
         mInjector.profileTabsController = new ProfileTabsController(
                 mInjector.selectionMgr,
diff --git a/src/com/android/documentsui/queries/SearchViewManager.java b/src/com/android/documentsui/queries/SearchViewManager.java
index 053dc93c8..1e26e4e08 100644
--- a/src/com/android/documentsui/queries/SearchViewManager.java
+++ b/src/com/android/documentsui/queries/SearchViewManager.java
@@ -20,6 +20,7 @@ import static com.android.documentsui.base.SharedMinimal.DEBUG;
 import static com.android.documentsui.base.State.ACTION_GET_CONTENT;
 import static com.android.documentsui.base.State.ACTION_OPEN;
 import static com.android.documentsui.base.State.ActionType;
+import static com.android.documentsui.flags.Flags.useMaterial3;
 
 import android.content.Intent;
 import android.os.Bundle;
@@ -332,7 +333,14 @@ public class SearchViewManager implements
         }
 
         // Recent root show open search bar, do not show duplicate search icon.
-        mMenuItem.setVisible(supportsSearch && (!stack.isRecents() || !mShowSearchBar));
+        boolean enabled = supportsSearch && (!stack.isRecents() || !mShowSearchBar);
+        mMenuItem.setVisible(enabled);
+        if (useMaterial3()) {
+            // When the use_material3 flag is enabled, we inflate and deflate the menu.
+            // This causes the search button to be disabled on inflation, toggle it in
+            // this scenario.
+            mMenuItem.setEnabled(enabled);
+        }
 
         mChipViewManager.setChipsRowVisible(supportsSearch && root.supportsMimeTypesSearch());
     }
diff --git a/tests/common/com/android/documentsui/bots/UiBot.java b/tests/common/com/android/documentsui/bots/UiBot.java
index f30cb93b8..4ec75bdc6 100644
--- a/tests/common/com/android/documentsui/bots/UiBot.java
+++ b/tests/common/com/android/documentsui/bots/UiBot.java
@@ -25,6 +25,8 @@ import static androidx.test.espresso.matcher.ViewMatchers.withClassName;
 import static androidx.test.espresso.matcher.ViewMatchers.withId;
 import static androidx.test.espresso.matcher.ViewMatchers.withText;
 
+import static com.android.documentsui.flags.Flags.useMaterial3;
+
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNotNull;
 import static junit.framework.Assert.assertNull;
@@ -68,37 +70,48 @@ import java.util.List;
  */
 public class UiBot extends Bots.BaseBot {
 
-    public static String targetPackageName;
-
     @SuppressWarnings("unchecked")
     private static final Matcher<View> TOOLBAR = allOf(
             isAssignableFrom(Toolbar.class),
             withId(R.id.toolbar));
-
     @SuppressWarnings("unchecked")
     private static final Matcher<View> ACTIONBAR = allOf(
             withClassName(endsWith("ActionBarContextView")));
-
     @SuppressWarnings("unchecked")
     private static final Matcher<View> TEXT_ENTRY = allOf(
             withClassName(endsWith("EditText")));
-
     @SuppressWarnings("unchecked")
     private static final Matcher<View> TOOLBAR_OVERFLOW = allOf(
             withClassName(endsWith("OverflowMenuButton")),
             ViewMatchers.isDescendantOfA(TOOLBAR));
-
     @SuppressWarnings("unchecked")
     private static final Matcher<View> ACTIONBAR_OVERFLOW = allOf(
             withClassName(endsWith("OverflowMenuButton")),
             ViewMatchers.isDescendantOfA(ACTIONBAR));
 
+    public static String targetPackageName;
+
     public UiBot(UiDevice device, Context context, int timeout) {
         super(device, context, timeout);
         targetPackageName =
                 InstrumentationRegistry.getInstrumentation().getTargetContext().getPackageName();
     }
 
+    private static Matcher<Object> withToolbarTitle(final Matcher<CharSequence> textMatcher) {
+        return new BoundedMatcher<Object, Toolbar>(Toolbar.class) {
+            @Override
+            public boolean matchesSafely(Toolbar toolbar) {
+                return textMatcher.matches(toolbar.getTitle());
+            }
+
+            @Override
+            public void describeTo(Description description) {
+                description.appendText("with toolbar title: ");
+                textMatcher.describeTo(description);
+            }
+        };
+    }
+
     public void assertWindowTitle(String expected) {
         onView(TOOLBAR)
                 .check(matches(withToolbarTitle(is(expected))));
@@ -198,7 +211,11 @@ public class UiBot extends Bots.BaseBot {
     }
 
     public void clickActionbarOverflowItem(String label) {
-        onView(ACTIONBAR_OVERFLOW).perform(click());
+        if (useMaterial3()) {
+            onView(TOOLBAR_OVERFLOW).perform(click());
+        } else {
+            onView(ACTIONBAR_OVERFLOW).perform(click());
+        }
         // Click the item by label, since Espresso doesn't support lookup by id on overflow.
         onView(withText(label)).perform(click());
     }
@@ -214,9 +231,10 @@ public class UiBot extends Bots.BaseBot {
     }
 
     public boolean waitForActionModeBarToAppear() {
+        String actionModeId = useMaterial3() ? "toolbar" : "action_mode_bar";
         UiObject2 bar =
-                mDevice.wait(Until.findObject(
-                        By.res(mTargetPackage + ":id/action_mode_bar")), mTimeout);
+                mDevice.wait(
+                        Until.findObject(By.res(mTargetPackage + ":id/" + actionModeId)), mTimeout);
         return (bar != null);
     }
 
@@ -307,20 +325,4 @@ public class UiBot extends Bots.BaseBot {
         // TODO: use the system string ? android.R.string.action_menu_overflow_description
         return mDevice.findObject(selector);
     }
-
-    private static Matcher<Object> withToolbarTitle(
-            final Matcher<CharSequence> textMatcher) {
-        return new BoundedMatcher<Object, Toolbar>(Toolbar.class) {
-            @Override
-            public boolean matchesSafely(Toolbar toolbar) {
-                return textMatcher.matches(toolbar.getTitle());
-            }
-
-            @Override
-            public void describeTo(Description description) {
-                description.appendText("with toolbar title: ");
-                textMatcher.describeTo(description);
-            }
-        };
-    }
 }
diff --git a/tests/unit/com/android/documentsui/files/ActionHandlerTest.java b/tests/unit/com/android/documentsui/files/ActionHandlerTest.java
index 3d9cba6bf..01dfa1c76 100644
--- a/tests/unit/com/android/documentsui/files/ActionHandlerTest.java
+++ b/tests/unit/com/android/documentsui/files/ActionHandlerTest.java
@@ -16,6 +16,7 @@
 
 package com.android.documentsui.files;
 
+import static com.android.documentsui.flags.Flags.useMaterial3;
 import static com.android.documentsui.testing.IntentAsserts.assertHasAction;
 import static com.android.documentsui.testing.IntentAsserts.assertHasData;
 import static com.android.documentsui.testing.IntentAsserts.assertHasExtra;
@@ -31,6 +32,8 @@ import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.Activity;
 import android.app.DownloadManager;
@@ -89,6 +92,8 @@ import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameter;
 import org.junit.runners.Parameterized.Parameters;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -108,6 +113,7 @@ public class ActionHandlerTest {
     private TestFeatures mFeatures;
     private TestConfigStore mTestConfigStore;
     private boolean refreshAnswer = false;
+    @Mock private Runnable mMockCloseSelectionBar;
 
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
@@ -126,6 +132,7 @@ public class ActionHandlerTest {
 
     @Before
     public void setUp() {
+        MockitoAnnotations.initMocks(this);
         mFeatures = new TestFeatures();
         mEnv = TestEnv.create(mFeatures);
         mActivity = TestActivity.create(mEnv);
@@ -152,6 +159,14 @@ public class ActionHandlerTest {
         mEnv.selectDocument(TestEnv.FILE_GIF);
     }
 
+    private void assertSelectionContainerClosed() {
+        if (useMaterial3()) {
+            verify(mMockCloseSelectionBar, times(1)).run();
+        } else {
+            assertTrue(mActionModeAddons.finishActionModeCalled);
+        }
+    }
+
     @Test
     public void testOpenSelectedInNewWindow() {
         mHandler.openSelectedInNewWindow();
@@ -195,7 +210,7 @@ public class ActionHandlerTest {
     @Test
     public void testSpringOpenDirectory() {
         mHandler.springOpenDirectory(TestEnv.FOLDER_0);
-        assertTrue(mActionModeAddons.finishActionModeCalled);
+        assertSelectionContainerClosed();
         assertEquals(TestEnv.FOLDER_0, mEnv.state.stack.peek());
     }
 
@@ -250,7 +265,7 @@ public class ActionHandlerTest {
         mHandler.deleteSelectedDocuments(docs, mEnv.state.stack.peek());
 
         mActivity.startService.assertCalled();
-        assertTrue(mActionModeAddons.finishActionModeCalled);
+        assertSelectionContainerClosed();
     }
 
     @Test
@@ -845,10 +860,10 @@ public class ActionHandlerTest {
                 mEnv.searchViewManager,
                 mEnv::lookupExecutor,
                 mActionModeAddons,
+                mMockCloseSelectionBar,
                 mClipper,
-                null,  // clip storage, not utilized unless we venture into *jumbo* clip territory.
+                null, // clip storage, not utilized unless we venture into *jumbo* clip territory.
                 mDragAndDropManager,
-                mEnv.injector
-        );
+                mEnv.injector);
     }
 }