Persist desk UUIDs

Persist desk UUIDs on full restore to fix issue where windows could be
restored to the wrong desk if the desks were reordered before full
restore.

Bug: b/247946400
Change-Id: Id0bbfb307b0c548db0fde42056838b68910153d6
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4225218
Reviewed-by: Ben Becker <benbecker@chromium.org>
Reviewed-by: ccameron chromium <ccameron@chromium.org>
Commit-Queue: Andrew Pantera <andp@chromium.org>
Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1146678}

ash/public/cpp/desk_template.cc

@@ -82,9 +82,9 @@ std::unique_ptr<DeskTemplate> DeskTemplate::Clone() const {
   return desk_template;
 }
 
-void DeskTemplate::SetDeskIndex(int desk_index) {
+void DeskTemplate::SetDeskUuid(base::Uuid desk_uuid) {
   DCHECK(desk_restore_data_);
-  desk_restore_data_->SetDeskIndex(desk_index);
+  desk_restore_data_->SetDeskUuid(desk_uuid);
 }
 
 std::string DeskTemplate::ToString() const {

ash/public/cpp/desk_template.h

@@ -133,9 +133,8 @@ class ASH_PUBLIC_EXPORT DeskTemplate {
   // Indicates whether this template can be modified by user.
   bool IsModifiable() const { return source_ == DeskTemplateSource::kUser; }
 
-  // Sets `desk_index` as the desk to launch on for all windows in the
-  // template.
-  void SetDeskIndex(int desk_index);
+  // Sets `desk_uuid` as the desk to launch on for all windows in the template.
+  void SetDeskUuid(base::Uuid desk_uuid);
 
   // Returns `this` in string format. Used for feedback logs.
   std::string ToString() const;

ash/public/cpp/window_properties.cc

@@ -67,6 +67,5 @@ DEFINE_OWNED_UI_CLASS_PROPERTY_KEY(gfx::Rect,
 DEFINE_UI_CLASS_PROPERTY_KEY(ResizeShadowType,
                              kResizeShadowTypeKey,
                              ResizeShadowType::kUnlock)
-DEFINE_OWNED_UI_CLASS_PROPERTY_KEY(std::string, kDeskGuidKey, nullptr)
 
 }  // namespace ash

ash/public/cpp/window_properties.h

@@ -194,11 +194,6 @@ ASH_PUBLIC_EXPORT extern const aura::WindowProperty<bool>* const
 ASH_PUBLIC_EXPORT extern const aura::WindowProperty<gfx::Rect*>* const
     kWindowPipResizeHandleBoundsKey;
 
-// A property key to indicate a desk guid of a workspace this window belongs
-// to.
-ASH_PUBLIC_EXPORT extern const aura::WindowProperty<std::string*>* const
-    kDeskGuidKey;
-
 // Alphabetical sort.
 
 }  // namespace ash

ash/wm/always_on_top_controller.cc

@@ -6,6 +6,7 @@
 
 #include "ash/public/cpp/shell_window_ids.h"
 #include "ash/public/cpp/window_properties.h"
+#include "ash/wm/desks/desk.h"
 #include "ash/wm/desks/desks_controller.h"
 #include "ash/wm/desks/desks_util.h"
 #include "ash/wm/window_state.h"
@@ -58,14 +59,28 @@ aura::Window* AlwaysOnTopController::GetContainer(aura::Window* window) const {
       ui::ZOrderLevel::kNormal) {
     aura::Window* root = always_on_top_container_->GetRootWindow();
 
-    // TODO(afakhry): Do we need to worry about the context of |window| here? Or
-    // is it safe to assume that |window| should always be parented to the
-    // active desks' container.
+    DesksController* desks_controller = DesksController::Get();
+    const std::string* desk_uuid_string =
+        window->GetProperty(aura::client::kDeskUuidKey);
+    if (desk_uuid_string) {
+      const base::Uuid desk_guid =
+          base::Uuid::ParseLowercase(*desk_uuid_string);
+      if (desk_guid.is_valid()) {
+        if (Desk* target_desk = desks_controller->GetDeskByUuid(desk_guid)) {
+          if (auto* container = target_desk->GetDeskContainerForRoot(root)) {
+            return container;
+          }
+        }
+      }
+    }
+
     const int window_workspace =
         window->GetProperty(aura::client::kWindowWorkspaceKey);
     if (window_workspace != aura::client::kWindowWorkspaceUnassignedWorkspace) {
+      LOG(ERROR) << "AlwaysOnTopController::GetContainer 2";
       auto* desk_container =
-          DesksController::Get()->GetDeskContainer(root, window_workspace);
+          desks_controller->GetDeskContainer(root, window_workspace);
+      LOG(ERROR) << "AlwaysOnTopController::GetContainer 3";
       if (desk_container)
         return desk_container;
     }

ash/wm/desks/desk.cc

@@ -396,9 +396,9 @@ void Desk::AddWindowToDesk(aura::Window* window) {
   if (!is_desk_being_removed_ &&
       !desks_util::IsWindowVisibleOnAllWorkspaces(window)) {
     // Setting the property for `kWindowWorkspaceKey` or
-    // `kDeskGuidKey` will trigger a save for the window state. To
+    // `kDeskUuidKey` will trigger a save for the window state. To
     // avoid doing this twice, we tell the window state to hold off on saving
-    // until we save the `kDeskGuidKey` value.
+    // until we save the `kDeskUuidKey` value.
     // TODO(b/265490703): We should eventually clean up this and
     // `GetScopedIgnorePropertyChange` when unit tests no longer need this
     // scoping to prevent double saves.
@@ -409,7 +409,7 @@ void Desk::AddWindowToDesk(aura::Window* window) {
                           desks_controller->GetDeskIndex(this));
     }
 
-    window->SetProperty(kDeskGuidKey, uuid_.AsLowercaseString());
+    window->SetProperty(aura::client::kDeskUuidKey, uuid_.AsLowercaseString());
   }
 
   MaybeIncrementWeeklyActiveDesks();

ash/wm/desks/desks_controller.cc

@@ -555,6 +555,15 @@ Desk* DesksController::GetDeskByUuid(const base::Uuid& desk_uuid) const {
   return it != desks_.end() ? it->get() : nullptr;
 }
 
+int DesksController::GetDeskIndexByUuid(const base::Uuid& desk_uuid) const {
+  for (size_t i = 0; i < desks_.size(); ++i) {
+    if (desk_uuid == desks_[i]->uuid()) {
+      return i;
+    }
+  }
+  return -1;
+}
+
 bool DesksController::CanRemoveDesks() const {
   return desks_.size() > 1;
 }
@@ -1276,13 +1285,11 @@ bool DesksController::OnSingleInstanceAppLaunchingFromSavedDesk(
   // not be visible on all desks.
   if (!desks_util::IsWindowVisibleOnAllWorkspaces(
           existing_app_instance_window)) {
-    // The index of the target desk is found in `app_restore_data`. If it isn't
-    // set, or out of bounds, then we default to the rightmost desk.
-    const int rightmost_desk_index = desks_.size() - 1;
-    const int target_desk_index =
-        std::min(app_restore_data.desk_id.value_or(rightmost_desk_index),
-                 rightmost_desk_index);
-    Desk* target_desk = desks_[target_desk_index].get();
+    // The uuid of the target desk is found in `app_restore_data`. If it isn't
+    // set, or is invalid, then we default to the rightmost desk.
+    Desk* target_desk = app_restore_data.desk_guid.is_valid()
+                            ? GetDeskByUuid(app_restore_data.desk_guid)
+                            : desks_.back().get();
 
     DCHECK(src_desk);
     if (src_desk != target_desk) {

ash/wm/desks/desks_controller.h

@@ -186,6 +186,10 @@ class ASH_EXPORT DesksController : public chromeos::DesksHelper,
   // found.
   Desk* GetDeskByUuid(const base::Uuid& desk_uuid) const;
 
+  // Returns the desk index of the desk that matches the desk_uuid, and returns
+  // -1 if no match is found.
+  int GetDeskIndexByUuid(const base::Uuid& desk_uuid) const;
+
   // Creates a new desk. CanCreateDesks() must be checked before calling this.
   void NewDesk(DesksCreationRemovalSource source);
 

ash/wm/desks/templates/admin_template_launch_tracker.cc

@@ -358,9 +358,7 @@ void AdminTemplateLaunchTracker::LaunchTemplate(SavedDeskDelegate* delegate,
 
   // Set apps to launch on the current desk.
   auto* desks_controller = DesksController::Get();
-  const int desk_index =
-      desks_controller->GetDeskIndex(desks_controller->active_desk());
-  admin_template->SetDeskIndex(desk_index);
+  admin_template->SetDeskUuid(desks_controller->active_desk()->uuid());
 
   UpdateAdminTemplateActivationIndices(*admin_template);
 

ash/wm/desks/templates/restore_data_collector.cc

@@ -79,9 +79,10 @@ void RestoreDataCollector::CaptureActiveDeskAsSavedDesk(
         BuildWindowInfo(window, /*activation_index=*/absl::nullopt,
                         /*for_saved_desks=*/true, mru_windows);
 
-    // Clear the desk ID in the WindowInfo that is to be stored in the template.
-    // It will be set to the ID of a newly created desk when launching.
+    // Clear the desk ID and uuid in the WindowInfo that is to be stored in the
+    // template. They will be set to the newly created desk when launching.
     window_info->desk_id.reset();
+    window_info->desk_guid = base::Uuid();
 
     ++call.pending_request_count;
     delegate->GetAppLaunchDataForSavedDesk(

ash/wm/desks/templates/saved_desk_presenter.cc

@@ -572,10 +572,9 @@ void SavedDeskPresenter::LaunchSavedDeskIntoNewDesk(
     }
   }
 
-  // Copy the index of the newly created desk to the saved desk. This ensures
+  // Copy the uuid of the newly created desk to the saved desk. This ensures
   // that apps appear on the right desk even if the user switches to another.
-  const int desk_index = DesksController::Get()->GetDeskIndex(new_desk);
-  saved_desk->SetDeskIndex(desk_index);
+  saved_desk->SetDeskUuid(new_desk->uuid());
 
   Shell::Get()->saved_desk_delegate()->LaunchAppsFromSavedDesk(
       std::move(saved_desk));
@@ -601,7 +600,7 @@ void SavedDeskPresenter::LaunchSavedDeskIntoNewDesk(
 
   overview_session_->GetGridWithRootWindow(root_window)
       ->desks_bar_view()
-      ->NudgeDeskName(desk_index);
+      ->NudgeDeskName(DesksController::Get()->GetDeskIndex(new_desk));
 
   if (saved_desk_type == DeskTemplateType::kSaveAndRecall) {
     // Passing nullopt as type since this indicates that we don't want to record

ash/wm/splitview/split_view_metrics_controller.cc

@@ -411,9 +411,8 @@ void SplitViewMetricsController::OnWindowInitialized(aura::Window* window) {
   }
 
   // Check if the recovered window is in the current desk.
-  if (!window_info->desk_id.has_value() ||
-      window_info->desk_id.value() !=
-          DesksController::Get()->GetDeskIndex(current_desk_)) {
+  if (!window_info->desk_guid.is_valid() ||
+      window_info->desk_guid != current_desk_->uuid()) {
     return;
   }
 

ash/wm/window_restore/window_restore_controller.cc

@@ -14,6 +14,7 @@
 #include "ash/session/session_controller_impl.h"
 #include "ash/shell.h"
 #include "ash/wm/container_finder.h"
+#include "ash/wm/desks/desk.h"
 #include "ash/wm/desks/desks_util.h"
 #include "ash/wm/desks/templates/saved_desk_util.h"
 #include "ash/wm/float/float_controller.h"
@@ -343,9 +344,18 @@ void WindowRestoreController::OnParentWindowToValidContainer(
 
   app_restore::WindowInfo* window_info = GetWindowInfo(window);
   if (window_info) {
-    const int desk_id = window_info->desk_id
-                            ? int{*window_info->desk_id}
-                            : aura::client::kWindowWorkspaceUnassignedWorkspace;
+    int desk_id = -1;
+    if (window_info->desk_guid.is_valid()) {
+      desk_id =
+          DesksController::Get()->GetDeskIndexByUuid(window_info->desk_guid);
+    }
+    // Its possible that the uuid is valid but it is not the uuid of any current
+    // desk.
+    if (desk_id == -1) {
+      desk_id = window_info->desk_id
+                    ? static_cast<int>(*window_info->desk_id)
+                    : aura::client::kWindowWorkspaceUnassignedWorkspace;
+    }
     window->SetProperty(aura::client::kWindowWorkspaceKey, desk_id);
   }
 

ash/wm/window_restore/window_restore_util.cc

@@ -62,12 +62,13 @@ std::unique_ptr<app_restore::WindowInfo> BuildWindowInfo(
     window_info->activation_index = window_activation_index;
   window_info->window = window;
   window_info->desk_id = window->GetProperty(aura::client::kWindowWorkspaceKey);
-  const std::string* desk_guid = window->GetProperty(kDeskGuidKey);
+  const std::string* desk_uuid =
+      window->GetProperty(aura::client::kDeskUuidKey);
 
   // It's possible for the desk to no longer exist or not be found in the case of
   // CloseAll.
   window_info->desk_guid =
-      desk_guid ? base::Uuid::ParseLowercase(*desk_guid) : base::Uuid();
+      desk_uuid ? base::Uuid::ParseLowercase(*desk_uuid) : base::Uuid();
 
   // If override bounds and window state are available (in tablet mode), save
   // those bounds.

ash/wm/window_state.cc

@@ -1376,7 +1376,8 @@ void WindowState::OnWindowPropertyChanged(aura::Window* window,
     }
     return;
   }
-  if (key == aura::client::kWindowWorkspaceKey || key == kDeskGuidKey) {
+  if (key == aura::client::kWindowWorkspaceKey ||
+      key == aura::client::kDeskUuidKey) {
     // Save the window for window restore purposes unless
     // |ignore_property_change_| is true. Note that moving windows across
     // displays will also trigger a kWindowWorkspaceKey change, even if the

chrome/browser/ash/app_restore/full_restore_app_launch_handler_browsertest.cc

@@ -106,6 +106,7 @@ constexpr char kRestoreIdPrefName[] = "browser_restore_id";
 // Test values for a test WindowInfo object.
 constexpr int kActivationIndex = 2;
 constexpr int kDeskId = 2;
+const base::Uuid kDeskUuid = base::Uuid::GenerateRandomV4();
 constexpr gfx::Rect kCurrentBounds(500, 200);
 constexpr chromeos::WindowStateType kWindowStateType =
     chromeos::WindowStateType::kPrimarySnapped;
@@ -165,6 +166,7 @@ class TestAppRestoreInfoObserver
 
 // Creates a WindowInfo object and then saves it.
 void CreateAndSaveWindowInfo(int desk_id,
+                             const base::Uuid& desk_uuid,
                              const gfx::Rect& current_bounds,
                              chromeos::WindowStateType window_state_type,
                              ui::WindowShowState pre_minimized_show_state,
@@ -180,6 +182,7 @@ void CreateAndSaveWindowInfo(int desk_id,
   ::app_restore::WindowInfo window_info;
   window_info.window = window.get();
   window_info.desk_id = desk_id;
+  window_info.desk_guid = desk_uuid;
   window_info.current_bounds = current_bounds;
   window_info.window_state_type = window_state_type;
 
@@ -202,6 +205,7 @@ void SaveWindowInfo(aura::Window* window) {
   window_info.window = window;
   window_info.activation_index = kActivationIndex;
   window_info.desk_id = kDeskId;
+  window_info.desk_guid = kDeskUuid;
   window_info.current_bounds = kCurrentBounds;
   window_info.window_state_type = WindowState::Get(window)->GetStateType();
   ::full_restore::SaveWindowInfo(window_info);
@@ -215,6 +219,7 @@ void SaveWindowInfo(
   window_info.window = window;
   window_info.activation_index = activation_index;
   window_info.desk_id = kDeskId;
+  window_info.desk_guid = kDeskUuid;
   window_info.current_bounds = kCurrentBounds;
   window_info.window_state_type = window_state_type;
   ::full_restore::SaveWindowInfo(window_info);
@@ -433,7 +438,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerBrowserTest,
           WindowOpenDisposition::NEW_WINDOW, display::kDefaultDisplayId,
           std::vector<base::FilePath>{}, nullptr));
   CreateAndSaveWindowInfo(
-      kDeskId, kCurrentBounds, chromeos::WindowStateType::kMinimized,
+      kDeskId, kDeskUuid, kCurrentBounds, chromeos::WindowStateType::kMinimized,
       ui::SHOW_STATE_MAXIMIZED, kWindowId2, /*snap_percentage=*/0);
 
   WaitForAppLaunchInfoSaved();
@@ -896,7 +901,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerBrowserTest,
                                 app_constants::kChromeAppId, kWindowId1));
 
   constexpr uint32_t kSnapPercentage = 75;
-  CreateAndSaveWindowInfo(kDeskId, kCurrentBounds, kWindowStateType,
+  CreateAndSaveWindowInfo(kDeskId, kDeskUuid, kCurrentBounds, kWindowStateType,
                           ui::SHOW_STATE_DEFAULT, kWindowId1, kSnapPercentage);
   WaitForAppLaunchInfoSaved();
 
@@ -914,6 +919,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerBrowserTest,
   window->SetProperty(::app_restore::kRestoreWindowIdKey, kWindowId1);
   auto stored_window_info = GetWindowInfo(window.get());
   EXPECT_EQ(kDeskId, *stored_window_info->desk_id);
+  EXPECT_EQ(kDeskUuid, stored_window_info->desk_guid);
   EXPECT_EQ(kCurrentBounds, *stored_window_info->current_bounds);
   EXPECT_EQ(kWindowStateType, *stored_window_info->window_state_type);
   EXPECT_EQ(kSnapPercentage, *stored_window_info->snap_percentage);
@@ -968,7 +974,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerBrowserTest,
       std::make_unique<::app_restore::AppLaunchInfo>(
           app_constants::kChromeAppId, previous_browser_id));
   CreateAndSaveWindowInfo(
-      kDeskId, kCurrentBounds, chromeos::WindowStateType::kNormal,
+      kDeskId, kDeskUuid, kCurrentBounds, chromeos::WindowStateType::kNormal,
       ui::SHOW_STATE_DEFAULT, previous_browser_id, /*snap_percentage=*/0);
   WaitForAppLaunchInfoSaved();
 
@@ -1001,7 +1007,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerBrowserTest,
       profile()->GetPath(), std::make_unique<::app_restore::AppLaunchInfo>(
                                 app_constants::kLacrosAppId, kWindowId1));
   CreateAndSaveWindowInfo(
-      kDeskId, prerestore_bounds, chromeos::WindowStateType::kNormal,
+      kDeskId, kDeskUuid, prerestore_bounds, chromeos::WindowStateType::kNormal,
       ui::SHOW_STATE_DEFAULT, kWindowId1, /*snap_percentage=*/0);
   WaitForAppLaunchInfoSaved();
 
@@ -1154,6 +1160,7 @@ IN_PROC_BROWSER_TEST_F(FullRestoreAppLaunchHandlerChromeAppBrowserTest,
   AutotestDesksApi().CreateNewDesk();
   AutotestDesksApi().CreateNewDesk();
   AutotestDesksApi().CreateNewDesk();
+  ActivateDesk(2);
 
   ::full_restore::SetActiveProfilePath(profile()->GetPath());
 

chrome/browser/ui/ash/desks/desks_client.cc

@@ -644,9 +644,9 @@ void DesksClient::OnGetTemplateForDeskLaunch(
     return;
   }
 
-  // Copy the index of the newly created desk to the saved desk. This ensures
+  // Copy the uuid of the newly created desk to the saved desk. This ensures
   // that apps appear on the right desk even if the user switches to another.
-  saved_desk->SetDeskIndex(desks_controller_->GetDeskIndex(new_desk));
+  saved_desk->SetDeskUuid(new_desk->uuid());
 
   const auto saved_desk_type = saved_desk->type();
   const auto uuid = saved_desk->uuid();

components/app_restore/app_restore_data.cc

@@ -32,7 +32,7 @@ constexpr char kAppNameKey[] = "app_name";
 constexpr char kActivationIndexKey[] = "index";
 constexpr char kFirstNonPinnedTabIndexKey[] = "first_non_pinned_tab_index";
 constexpr char kDeskIdKey[] = "desk_id";
-constexpr char kDeskGuidKey[] = "desk_guid";
+constexpr char kDeskUuidKey[] = "desk_guid";
 constexpr char kCurrentBoundsKey[] = "current_bounds";
 constexpr char kWindowStateTypeKey[] = "window_state_type";
 constexpr char kPreMinimizedShowStateTypeKey[] = "pre_min_state";
@@ -247,7 +247,7 @@ AppRestoreData::AppRestoreData(base::Value::Dict&& data) {
   first_non_pinned_tab_index =
       GetIntValueFromDict(data, kFirstNonPinnedTabIndexKey);
   desk_id = GetIntValueFromDict(data, kDeskIdKey);
-  desk_guid = GetGuidValueFromDict(data, kDeskGuidKey);
+  desk_guid = GetGuidValueFromDict(data, kDeskUuidKey);
   current_bounds = GetBoundsRectFromDict(data, kCurrentBoundsKey);
   window_state_type = GetWindowStateTypeFromDict(data);
   pre_minimized_show_state_type = GetPreMinimizedShowStateTypeFromDict(data);
@@ -443,7 +443,7 @@ base::Value AppRestoreData::ConvertToValue() const {
     launch_info_dict.Set(kDeskIdKey, desk_id.value());
 
   if (desk_guid.is_valid()) {
-    launch_info_dict.Set(kDeskGuidKey, desk_guid.AsLowercaseString());
+    launch_info_dict.Set(kDeskUuidKey, desk_guid.AsLowercaseString());
   }
 
   if (current_bounds.has_value()) {