Add crostini helper to handle Crostini app window special cases.

AppServiceAppWindowCrostiniTracker is added to handle Crostini app
window special cases, e.g. CrostiniAppDisplay, Crostini shelf app id,
etc.

AppServiceAppWindowCrostiniHelper implements part of functions from
CrostiniAppWindowShelfController.

Update Crostini's LaunchApplication and CrostiniAppRestartView::Accept
to use AppServiceAppWindowCrostiniHelper.

BUG=1011235

Change-Id: If8f71b7b45e190431b1023fe9f1feab746f788f4
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1941539
Commit-Queue: Nancy Wang <nancylingwang@chromium.org>
Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Reviewed-by: Joel Hockey <joelhockey@chromium.org>
Cr-Commit-Position: refs/heads/master@{#721283}

chrome/browser/chromeos/crostini/crostini_util.cc

@@ -29,6 +29,8 @@
 #include "chrome/browser/chromeos/virtual_machines/virtual_machines_util.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/app_list/crostini/crostini_app_icon.h"
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.h"
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_launcher_controller.h"
 #include "chrome/browser/ui/ash/launcher/chrome_launcher_controller.h"
 #include "chrome/browser/ui/ash/launcher/crostini_app_window_shelf_controller.h"
 #include "chrome/browser/ui/ash/launcher/shelf_spinner_controller.h"
@@ -162,10 +164,19 @@ void LaunchApplication(
   ChromeLauncherController* chrome_launcher_controller =
       ChromeLauncherController::instance();
   DCHECK(chrome_launcher_controller);
-  CrostiniAppWindowShelfController* shelf_controller =
-      chrome_launcher_controller->crostini_app_window_shelf_controller();
-  DCHECK(shelf_controller);
-  shelf_controller->OnAppLaunchRequested(app_id, display_id);
+
+  if (base::FeatureList::IsEnabled(features::kAppServiceInstanceRegistry)) {
+    AppServiceAppWindowLauncherController* app_service_controller =
+        chrome_launcher_controller->app_service_app_window_controller();
+    DCHECK(app_service_controller);
+    app_service_controller->app_service_crostini_tracker()
+        ->OnAppLaunchRequested(app_id, display_id);
+  } else {
+    CrostiniAppWindowShelfController* shelf_controller =
+        chrome_launcher_controller->crostini_app_window_shelf_controller();
+    DCHECK(shelf_controller);
+    shelf_controller->OnAppLaunchRequested(app_id, display_id);
+  }
 
   // Share any paths not in crostini.  The user will see the spinner while this
   // is happening.

chrome/browser/ui/BUILD.gn

@@ -3725,6 +3725,8 @@ jumbo_static_library("ui") {
       "ash/assistant/proactive_suggestions_client_impl.h",
       "ash/assistant/proactive_suggestions_loader.cc",
       "ash/assistant/proactive_suggestions_loader.h",
+      "ash/launcher/app_service_app_window_crostini_tracker.cc",
+      "ash/launcher/app_service_app_window_crostini_tracker.h",
       "ash/launcher/app_service_app_window_launcher_controller.cc",
       "ash/launcher/app_service_app_window_launcher_controller.h",
       "ash/launcher/app_service_instance_registry_helper.cc",

chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.h"
+
+#include "ash/public/cpp/multi_user_window_manager.h"
+#include "chrome/browser/chromeos/crostini/crostini_force_close_watcher.h"
+#include "chrome/browser/chromeos/crostini/crostini_registry_service.h"
+#include "chrome/browser/chromeos/crostini/crostini_registry_service_factory.h"
+#include "chrome/browser/chromeos/crostini/crostini_util.h"
+#include "chrome/browser/chromeos/plugin_vm/plugin_vm_util.h"
+#include "chrome/browser/chromeos/profiles/profile_helper.h"
+#include "chrome/browser/profiles/profile.h"
+#include "chrome/browser/ui/ash/launcher/chrome_launcher_controller.h"
+#include "chrome/browser/ui/ash/launcher/shelf_spinner_controller.h"
+#include "chrome/browser/ui/ash/multi_user/multi_user_window_manager_helper.h"
+#include "chrome/browser/ui/browser.h"
+#include "chrome/browser/ui/browser_finder.h"
+#include "chrome/common/chrome_features.h"
+#include "components/arc/arc_util.h"
+#include "components/exo/shell_surface_util.h"
+#include "components/user_manager/user_manager.h"
+#include "ui/aura/window.h"
+#include "ui/display/display.h"
+#include "ui/display/screen.h"
+#include "ui/wm/core/window_util.h"
+
+namespace {
+
+void MoveWindowFromOldDisplayToNewDisplay(aura::Window* window,
+                                          display::Display& old_display,
+                                          display::Display& new_display) {
+  // Adjust the window size and origin in proportion to the relative size of the
+  // display.
+  int old_width = old_display.bounds().width();
+  int new_width = new_display.bounds().width();
+  int old_height = old_display.bounds().height();
+  int new_height = new_display.bounds().height();
+  gfx::Rect old_bounds = window->bounds();
+  gfx::Rect new_bounds(old_bounds.x() * new_width / old_width,
+                       old_bounds.y() * new_height / old_height,
+                       old_bounds.width() * new_width / old_width,
+                       old_bounds.height() * new_height / old_height);
+
+  // Transform the bounds in display to that in screen.
+  gfx::Point new_origin = new_display.bounds().origin();
+  new_origin.Offset(new_bounds.x(), new_bounds.y());
+  new_bounds.set_origin(new_origin);
+  window->SetBoundsInScreen(new_bounds, new_display);
+}
+
+}  // namespace
+
+AppServiceAppWindowCrostiniTracker::AppServiceAppWindowCrostiniTracker() =
+    default;
+
+AppServiceAppWindowCrostiniTracker::~AppServiceAppWindowCrostiniTracker() =
+    default;
+
+void AppServiceAppWindowCrostiniTracker::OnWindowVisibilityChanging(
+    aura::Window* window,
+    const std::string& shelf_app_id) {
+  // Transient windows are set up after window init, so remove them here.
+  // Crostini shouldn't need to know about ARC app windows.
+  if (wm::GetTransientParent(window) ||
+      arc::GetWindowTaskId(window) != arc::kNoTaskId ||
+      plugin_vm::IsPluginVmWindow(window)) {
+    return;
+  }
+
+  // Handle browser windows, such as the Crostini terminal.
+  Browser* browser = chrome::FindBrowserWithWindow(window);
+  if (browser)
+    return;
+
+  // Currently Crostini can only be used from the primary profile. In the
+  // future, this may be replaced by some way of matching the container that
+  // runs this app with the user that owns it.
+  const AccountId& primary_account_id =
+      user_manager::UserManager::Get()->GetPrimaryUser()->GetAccountId();
+
+  crostini::CrostiniRegistryService* registry_service =
+      crostini::CrostiniRegistryServiceFactory::GetForProfile(
+          chromeos::ProfileHelper::Get()->GetProfileByAccountId(
+              primary_account_id));
+
+  // At this point, all remaining windows are Crostini windows. Firstly, we add
+  // support for forcibly closing it. We use the registration to retrieve the
+  // app's name, but this may be null in the case of apps with no associated
+  // launcher entry (i.e. no .desktop file), in which case the app's name is
+  // unknown.
+  base::Optional<crostini::CrostiniRegistryService::Registration> registration =
+      registry_service->GetRegistration(shelf_app_id);
+  RegisterCrostiniWindowForForceClose(
+      window, registration.has_value() ? registration->Name() : "");
+
+  // Failed to uniquely identify the Crostini app that this window is for.
+  // The spinners on the shelf have internal app IDs which are valid
+  // extensions IDs. If the ID here starts with "crostini:" then it implies
+  // that it has failed to identify the exact app that's starting.
+  // The existing spinner that fails to be linked back should be closed,
+  // otherwise it will be left on the shelf indefinetely until it is closed
+  // manually by the user.
+  // When the condition is triggered here, the container is up and at least
+  // one app is starting. It's safe to close all the spinners since their
+  // respective apps take at most another few seconds to start.
+  // Work is ongoing to make this occur as infrequently as possible.
+  // See https://crbug.com/854911.
+  if (base::StartsWith(shelf_app_id, crostini::kCrostiniAppIdPrefix,
+                       base::CompareCase::SENSITIVE)) {
+    ChromeLauncherController::instance()
+        ->GetShelfSpinnerController()
+        ->CloseCrostiniSpinners();
+  }
+
+  // Prevent Crostini window from showing up after user switch.
+  MultiUserWindowManagerHelper::GetWindowManager()->SetWindowOwner(
+      window, primary_account_id);
+
+  // Move the Crostini app window to the right display if necessary.
+  int64_t display_id = crostini_app_display_.GetDisplayIdForAppId(shelf_app_id);
+  if (display_id == display::kInvalidDisplayId)
+    return;
+
+  display::Display new_display;
+  if (!display::Screen::GetScreen()->GetDisplayWithDisplayId(display_id,
+                                                             &new_display)) {
+    return;
+  }
+
+  display::Display old_display =
+      display::Screen::GetScreen()->GetDisplayNearestWindow(window);
+
+  if (new_display != old_display)
+    MoveWindowFromOldDisplayToNewDisplay(window, old_display, new_display);
+}
+
+void AppServiceAppWindowCrostiniTracker::OnWindowDestroying(
+    const std::string& app_id) {
+  if (app_id != app_id_to_restart_)
+    return;
+  crostini::LaunchCrostiniApp(ChromeLauncherController::instance()->profile(),
+                              app_id, display_id_to_restart_in_);
+  app_id_to_restart_.clear();
+}
+
+void AppServiceAppWindowCrostiniTracker::OnAppLaunchRequested(
+    const std::string& app_id,
+    int64_t display_id) {
+  crostini_app_display_.Register(app_id, display_id);
+}
+
+void AppServiceAppWindowCrostiniTracker::Restart(const ash::ShelfID& shelf_id,
+                                                 int64_t display_id) {
+  app_id_to_restart_ = shelf_id.app_id;
+  display_id_to_restart_in_ = display_id;
+  ChromeLauncherController::instance()->Close(shelf_id);
+}
+
+std::string AppServiceAppWindowCrostiniTracker::GetShelfAppId(
+    aura::Window* window) const {
+  // Transient windows are set up after window init, so remove them here.
+  // Crostini shouldn't need to know about ARC app windows.
+  if (wm::GetTransientParent(window) ||
+      arc::GetWindowTaskId(window) != arc::kNoTaskId ||
+      plugin_vm::IsPluginVmWindow(window)) {
+    return std::string();
+  }
+
+  // Handle browser windows, such as the Crostini terminal.
+  Browser* browser = chrome::FindBrowserWithWindow(window);
+  if (browser) {
+    base::Optional<std::string> app_id =
+        crostini::CrostiniAppIdFromAppName(browser->app_name());
+    if (!app_id)
+      return std::string();
+    return app_id.value();
+  }
+
+  // Currently Crostini can only be used from the primary profile. In the
+  // future, this may be replaced by some way of matching the container that
+  // runs this app with the user that owns it.
+  const AccountId& primary_account_id =
+      user_manager::UserManager::Get()->GetPrimaryUser()->GetAccountId();
+  crostini::CrostiniRegistryService* registry_service =
+      crostini::CrostiniRegistryServiceFactory::GetForProfile(
+          chromeos::ProfileHelper::Get()->GetProfileByAccountId(
+              primary_account_id));
+  std::string shelf_app_id = registry_service->GetCrostiniShelfAppId(
+      exo::GetShellApplicationId(window), exo::GetShellStartupId(window));
+  return shelf_app_id;
+}
+
+void AppServiceAppWindowCrostiniTracker::RegisterCrostiniWindowForForceClose(
+    aura::Window* window,
+    const std::string& app_name) {
+  if (!base::FeatureList::IsEnabled(features::kCrostiniForceClose))
+    return;
+  exo::ShellSurfaceBase* surface = exo::GetShellSurfaceBaseForWindow(window);
+  if (!surface)
+    return;
+  crostini::ForceCloseWatcher::Watch(
+      std::make_unique<crostini::ShellSurfaceForceCloseDelegate>(surface,
+                                                                 app_name));
+}

chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_UI_ASH_LAUNCHER_APP_SERVICE_APP_WINDOW_CROSTINI_TRACKER_H_
+#define CHROME_BROWSER_UI_ASH_LAUNCHER_APP_SERVICE_APP_WINDOW_CROSTINI_TRACKER_H_
+
+#include "ash/public/cpp/shelf_types.h"
+#include "chrome/browser/ui/ash/launcher/crostini_app_display.h"
+
+namespace aura {
+class Window;
+}
+
+// AppServiceAppWindowCrostiniTracker is used to handle Crostini app window
+// special cases, e.g. CrostiniAppDisplay, Crostini shelf app id, etc.
+class AppServiceAppWindowCrostiniTracker {
+ public:
+  AppServiceAppWindowCrostiniTracker();
+  ~AppServiceAppWindowCrostiniTracker();
+
+  AppServiceAppWindowCrostiniTracker(
+      const AppServiceAppWindowCrostiniTracker&) = delete;
+  AppServiceAppWindowCrostiniTracker& operator=(
+      const AppServiceAppWindowCrostiniTracker&) = delete;
+
+  void OnWindowVisibilityChanging(aura::Window* window,
+                                  const std::string& shelf_app_id);
+  void OnWindowDestroying(const std::string& app_id);
+
+  // A Crostini app with |app_id| is requested to launch on display with
+  // |display_id|.
+  void OnAppLaunchRequested(const std::string& app_id, int64_t display_id);
+
+  // Close app with |shelf_id| and then restart it on |display_id|.
+  void Restart(const ash::ShelfID& shelf_id, int64_t display_id);
+
+  std::string GetShelfAppId(aura::Window* window) const;
+
+ private:
+  void RegisterCrostiniWindowForForceClose(aura::Window* window,
+                                           const std::string& app_name);
+
+  CrostiniAppDisplay crostini_app_display_;
+
+  // These two member variables track an app restart request. When
+  // app_id_to_restart_ is not empty the controller observes that app and
+  // relaunches it when all of its windows are closed.
+  std::string app_id_to_restart_;
+  int64_t display_id_to_restart_in_ = 0;
+};
+
+#endif  // CHROME_BROWSER_UI_ASH_LAUNCHER_APP_SERVICE_APP_WINDOW_CROSTINI_TRACKER_H_

chrome/browser/ui/ash/launcher/app_service_app_window_launcher_controller.cc

@@ -10,7 +10,7 @@
 #include "chrome/browser/apps/app_service/app_service_proxy_factory.h"
 #include "chrome/browser/chromeos/plugin_vm/plugin_vm_util.h"
 #include "chrome/browser/profiles/profile.h"
-#include "chrome/browser/profiles/profile_manager.h"
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.h"
 #include "chrome/browser/ui/ash/launcher/app_window_base.h"
 #include "chrome/browser/ui/ash/launcher/app_window_launcher_item_controller.h"
 #include "chrome/browser/ui/ash/launcher/chrome_launcher_controller.h"
@@ -29,11 +29,13 @@ AppServiceAppWindowLauncherController::AppServiceAppWindowLauncherController(
     : AppWindowLauncherController(owner),
       proxy_(apps::AppServiceProxyFactory::GetForProfile(owner->profile())),
       app_service_instance_helper_(
-          std::make_unique<AppServiceInstanceRegistryHelper>(
-              owner->profile())) {
+          std::make_unique<AppServiceInstanceRegistryHelper>(owner->profile())),
+      crostini_tracker_(
+          std::make_unique<AppServiceAppWindowCrostiniTracker>()) {
   aura::Env::GetInstance()->AddObserver(this);
   DCHECK(proxy_);
   DCHECK(app_service_instance_helper_);
+  DCHECK(crostini_tracker_);
   Observe(&proxy_->InstanceRegistry());
 }
 
@@ -103,29 +105,9 @@ void AppServiceAppWindowLauncherController::OnWindowVisibilityChanging(
   if (!observed_windows_.IsObserving(window))
     return;
 
-  ash::ShelfID shelf_id;
-  if (!proxy_->InstanceRegistry().ForOneInstance(
-          window, [&shelf_id](const apps::InstanceUpdate& update) {
-            shelf_id = ash::ShelfID(update.AppId(), update.LaunchId());
-          })) {
-    shelf_id = ash::ShelfID::Deserialize(window->GetProperty(ash::kShelfIDKey));
-  }
-
-  std::string app_id, launch_id;
-
-  if (shelf_id.IsNull()) {
-    if (!plugin_vm::IsPluginVmWindow(window))
-      return;
-    app_id = plugin_vm::kPluginVmAppId;
-    shelf_id = ash::ShelfID(plugin_vm::kPluginVmAppId);
-  } else {
-    if (proxy_->AppRegistryCache().GetAppType(shelf_id.app_id) ==
-        apps::mojom::AppType::kUnknown) {
-      return;
-    }
-    app_id = shelf_id.app_id;
-    launch_id = shelf_id.launch_id;
-  }
+  ash::ShelfID shelf_id = GetShelfId(window);
+  if (shelf_id.IsNull())
+    return;
 
   // Update |state|. The app must be started, and running state. If visible,
   // set it as |kVisible|, otherwise, clear the visible bit.
@@ -140,12 +122,15 @@ void AppServiceAppWindowLauncherController::OnWindowVisibilityChanging(
                     : static_cast<apps::InstanceState>(
                           state & ~(apps::InstanceState::kVisible));
 
-  app_service_instance_helper_->OnInstances(app_id, window, launch_id, state);
+  app_service_instance_helper_->OnInstances(shelf_id.app_id, window,
+                                            shelf_id.launch_id, state);
 
   if (!visible || shelf_id.app_id == extension_misc::kChromeAppId)
     return;
 
   RegisterAppWindow(window, shelf_id);
+
+  crostini_tracker_->OnWindowVisibilityChanging(window, shelf_id.app_id);
 }
 
 void AppServiceAppWindowLauncherController::OnWindowDestroying(
@@ -165,6 +150,9 @@ void AppServiceAppWindowLauncherController::OnWindowDestroying(
 
   RemoveFromShelf(app_window_it->second.get());
 
+  if (!shelf_id.IsNull())
+    crostini_tracker_->OnWindowDestroying(shelf_id.app_id);
+
   aura_window_to_app_window_.erase(app_window_it);
 }
 
@@ -377,3 +365,36 @@ bool AppServiceAppWindowLauncherController::IsOpenedInBrowserTab(
   }
   return false;
 }
+
+ash::ShelfID AppServiceAppWindowLauncherController::GetShelfId(
+    aura::Window* window) const {
+  std::string shelf_app_id = crostini_tracker_->GetShelfAppId(window);
+  if (!shelf_app_id.empty())
+    return ash::ShelfID(shelf_app_id);
+
+  ash::ShelfID shelf_id;
+
+  // If the window exists in InstanceRegistry, get the shelf id from
+  // InstanceRegistry.
+  bool exist_in_instance = proxy_->InstanceRegistry().ForOneInstance(
+      window, [&shelf_id](const apps::InstanceUpdate& update) {
+        shelf_id = ash::ShelfID(update.AppId(), update.LaunchId());
+      });
+  if (!exist_in_instance) {
+    shelf_id = ash::ShelfID::Deserialize(window->GetProperty(ash::kShelfIDKey));
+  }
+
+  if (!shelf_id.IsNull()) {
+    if (proxy_->AppRegistryCache().GetAppType(shelf_id.app_id) ==
+        apps::mojom::AppType::kUnknown) {
+      return ash::ShelfID();
+    }
+    return shelf_id;
+  }
+
+  // For null shelf id, it could be VM window or ARC apps window.
+  if (plugin_vm::IsPluginVmWindow(window))
+    return ash::ShelfID(plugin_vm::kPluginVmAppId);
+
+  return ash::ShelfID();
+}

chrome/browser/ui/ash/launcher/app_service_app_window_launcher_controller.h

@@ -23,6 +23,7 @@ namespace apps {
 class InstanceUpdate;
 }  // namespace apps
 
+class AppServiceAppWindowCrostiniTracker;
 class AppWindowBase;
 class ChromeLauncherController;
 
@@ -64,6 +65,10 @@ class AppServiceAppWindowLauncherController
   void OnInstanceRegistryWillBeDestroyed(
       apps::InstanceRegistry* instance_registry) override;
 
+  AppServiceAppWindowCrostiniTracker* app_service_crostini_tracker() {
+    return crostini_tracker_.get();
+  }
+
  private:
   using AuraWindowToAppWindow =
       std::map<aura::Window*, std::unique_ptr<AppWindowBase>>;
@@ -88,6 +93,8 @@ class AppServiceAppWindowLauncherController
   // Return true if the app is opened in a browser tab.
   bool IsOpenedInBrowserTab(const std::string& app_id);
 
+  ash::ShelfID GetShelfId(aura::Window* window) const;
+
   AuraWindowToAppWindow aura_window_to_app_window_;
   ScopedObserver<aura::Window, aura::WindowObserver> observed_windows_{this};
 
@@ -95,6 +102,8 @@ class AppServiceAppWindowLauncherController
   std::unique_ptr<AppServiceInstanceRegistryHelper>
       app_service_instance_helper_;
 
+  std::unique_ptr<AppServiceAppWindowCrostiniTracker> crostini_tracker_;
+
   DISALLOW_COPY_AND_ASSIGN(AppServiceAppWindowLauncherController);
 };
 

chrome/browser/ui/ash/launcher/chrome_launcher_controller.cc

@@ -270,8 +270,11 @@ ChromeLauncherController::ChromeLauncherController(Profile* profile,
       browser_status_monitor_ = std::make_unique<BrowserStatusMonitor>(this);
       browser_status_monitor_->Initialize();
     }
-    app_window_controllers_.push_back(
-        std::make_unique<AppServiceAppWindowLauncherController>(this));
+    std::unique_ptr<AppServiceAppWindowLauncherController>
+        app_service_controller =
+            std::make_unique<AppServiceAppWindowLauncherController>(this);
+    app_service_app_window_controller_ = app_service_controller.get();
+    app_window_controllers_.emplace_back(std::move(app_service_controller));
     return;
   }
 

chrome/browser/ui/ash/launcher/chrome_launcher_controller.h

@@ -26,6 +26,7 @@
 #include "components/sync_preferences/pref_service_syncable_observer.h"
 
 class AppIconLoader;
+class AppServiceAppWindowLauncherController;
 class AppWindowLauncherController;
 class BrowserShortcutLauncherItemController;
 class BrowserStatusMonitor;
@@ -78,6 +79,11 @@ class ChromeLauncherController
 
   Profile* profile() const { return profile_; }
   ash::ShelfModel* shelf_model() const { return model_; }
+
+  AppServiceAppWindowLauncherController* app_service_app_window_controller() {
+    return app_service_app_window_controller_;
+  }
+
   CrostiniAppWindowShelfController* crostini_app_window_shelf_controller()
       const {
     return crostini_app_window_shelf_controller_;
@@ -393,6 +399,10 @@ class ChromeLauncherController
   // The ShelfModel instance owned by ash::Shell's ShelfController.
   ash::ShelfModel* model_;
 
+  // The AppService app window launcher controller.
+  AppServiceAppWindowLauncherController* app_service_app_window_controller_ =
+      nullptr;
+
   // The shelf controller for Crostini apps.
   CrostiniAppWindowShelfController* crostini_app_window_shelf_controller_ =
       nullptr;

chrome/browser/ui/views/crostini/crostini_app_restart_view.cc

@@ -4,10 +4,13 @@
 
 #include "chrome/browser/ui/views/crostini/crostini_app_restart_view.h"
 
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_crostini_tracker.h"
+#include "chrome/browser/ui/ash/launcher/app_service_app_window_launcher_controller.h"
 #include "chrome/browser/ui/ash/launcher/chrome_launcher_controller.h"
 #include "chrome/browser/ui/ash/launcher/crostini_app_window_shelf_controller.h"
 #include "chrome/browser/ui/browser_dialogs.h"
 #include "chrome/browser/ui/views/chrome_layout_provider.h"
+#include "chrome/common/chrome_features.h"
 #include "ui/base/l10n/l10n_util.h"
 #include "ui/chromeos/devicetype_utils.h"
 #include "ui/display/screen.h"
@@ -45,9 +48,16 @@ bool CrostiniAppRestartView::ShouldShowCloseButton() const {
 }
 
 bool CrostiniAppRestartView::Accept() {
-  ChromeLauncherController::instance()
-      ->crostini_app_window_shelf_controller()
-      ->Restart(id_, display_id_);
+  if (base::FeatureList::IsEnabled(features::kAppServiceInstanceRegistry)) {
+    ChromeLauncherController::instance()
+        ->app_service_app_window_controller()
+        ->app_service_crostini_tracker()
+        ->Restart(id_, display_id_);
+  } else {
+    ChromeLauncherController::instance()
+        ->crostini_app_window_shelf_controller()
+        ->Restart(id_, display_id_);
+  }
   return true;  // Should close the dialog
 }