assistant: take screenshot on screen with stylus selection.

Bug: b/111105039
Test: Manual
Change-Id: I7c39e63c2b65f557f668687bdc34ced2256462f6
Reviewed-on: https://chromium-review.googlesource.com/1149065Reviewed-by: Mitsuru Oshima <oshima@chromium.org>
Reviewed-by: Xiaohui Chen <xiaohuic@chromium.org>
Commit-Queue: Muyuan Li <muyuanli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#579523}

ash/assistant/assistant_screen_context_controller.cc

@@ -21,6 +21,7 @@
 #include "ui/gfx/skbitmap_operations.h"
 #include "ui/snapshot/snapshot.h"
 #include "ui/snapshot/snapshot_aura.h"
+#include "ui/wm/core/focus_controller.h"
 #include "ui/wm/core/window_util.h"
 
 namespace ash {
@@ -143,8 +144,7 @@ void AssistantScreenContextController::RemoveModelObserver(
 void AssistantScreenContextController::RequestScreenshot(
     const gfx::Rect& rect,
     mojom::AssistantController::RequestScreenshotCallback callback) {
-  // TODO(muyuanli): Handle multi-display when Assistant's behavior is defined.
-  aura::Window* root_window = Shell::GetPrimaryRootWindow();
+  aura::Window* root_window = Shell::Get()->GetRootWindowForNewWindows();
 
   std::unique_ptr<ui::LayerTreeOwner> layer_owner =
       CreateLayerForAssistantSnapshot(root_window);

ash/highlighter/highlighter_controller.cc

@@ -13,6 +13,7 @@
 #include "ash/public/cpp/scale_utility.h"
 #include "ash/public/cpp/shell_window_ids.h"
 #include "ash/shell.h"
+#include "ash/shell_state.h"
 #include "ash/system/palette/palette_utils.h"
 #include "base/metrics/histogram_macros.h"
 #include "chromeos/chromeos_switches.h"
@@ -223,6 +224,10 @@ void HighlighterController::RecognizeGesture() {
 
   if (!box.IsEmpty() &&
       gesture_type != HighlighterGestureType::kNotRecognized) {
+    // The window for selection should be the root window to show assistant.
+    Shell::Get()->shell_state()->SetRootWindowForNewWindows(
+        current_window->GetRootWindow());
+
     // TODO(muyuanli): Delete the check when native assistant is default on.
     // This is a temporary workaround to support both ARC-based assistant
     // and native assistant. In ARC-based assistant, we send the rect in pixels