Fix screen context logic.

New logic:
- Cache screen context on UI launch.
- Start screen context interaction using cached data on chip press.
- Start metalayer (screenshot only) interaction with stylus.

Bug: b:113069944
Change-Id: Iea4089ac2411a8000bc7ebaa16e90e12fcb6a151
Reviewed-on: https://chromium-review.googlesource.com/1186160
Commit-Queue: David Black <dmblack@google.com>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Xiaohui Chen <xiaohuic@chromium.org>
Cr-Commit-Position: refs/heads/master@{#585992}

ash/assistant/assistant_interaction_controller.cc

@@ -72,8 +72,7 @@ void AssistantInteractionController::OnDeepLinkReceived(
   using namespace assistant::util;
 
   if (type == DeepLinkType::kWhatsOnMyScreen) {
-    // Start a screen context interaction using the entire screen region.
-    StartScreenContextInteraction(/*region=*/gfx::Rect());
+    StartScreenContextInteraction();
     return;
   }
 
@@ -137,8 +136,7 @@ void AssistantInteractionController::OnHighlighterEnabledChanged(
 
 void AssistantInteractionController::OnHighlighterSelectionRecognized(
     const gfx::Rect& rect) {
-  // Start a screen context interaction using the selected region.
-  StartScreenContextInteraction(rect);
+  StartMetalayerInteraction(/*region=*/rect);
 }
 
 void AssistantInteractionController::OnInteractionStateChanged(
@@ -380,7 +378,7 @@ void AssistantInteractionController::OnDialogPlateContentsCommitted(
   StartTextInteraction(text);
 }
 
-void AssistantInteractionController::StartScreenContextInteraction(
+void AssistantInteractionController::StartMetalayerInteraction(
     const gfx::Rect& region) {
   StopActiveInteraction();
 
@@ -388,13 +386,18 @@ void AssistantInteractionController::StartScreenContextInteraction(
       std::make_unique<AssistantTextQuery>(
           l10n_util::GetStringUTF8(IDS_ASH_ASSISTANT_CHIP_WHATS_ON_MY_SCREEN)));
 
-  // Screen context interactions are explicitly started by the user, either via
-  // suggestion chip or stylus selection, so we initiate a new request to ensure
-  // we have the freshest context possible. Because we specify |from_user| as
-  // being true, an interaction will be started to return contextual results
-  // from the server.
-  assistant_controller_->screen_context_controller()->RequestScreenContext(
-      region, /*from_user=*/true);
+  assistant_->StartMetalayerInteraction(region);
+}
+
+void AssistantInteractionController::StartScreenContextInteraction() {
+  StopActiveInteraction();
+
+  assistant_interaction_model_.SetPendingQuery(
+      std::make_unique<AssistantTextQuery>(
+          l10n_util::GetStringUTF8(IDS_ASH_ASSISTANT_CHIP_WHATS_ON_MY_SCREEN)));
+
+  // Note that screen context was cached when the UI was launched.
+  assistant_->StartCachedScreenContextInteraction();
 }
 
 void AssistantInteractionController::StartTextInteraction(

ash/assistant/assistant_interaction_controller.h

@@ -98,7 +98,8 @@ class AssistantInteractionController
   void OnSuggestionChipPressed(const AssistantSuggestion* suggestion);
 
  private:
-  void StartScreenContextInteraction(const gfx::Rect& region);
+  void StartMetalayerInteraction(const gfx::Rect& region);
+  void StartScreenContextInteraction();
   void StartTextInteraction(const std::string text);
   void StartVoiceInteraction();
   void StopActiveInteraction();

ash/assistant/assistant_screen_context_controller.cc

@@ -179,21 +179,6 @@ void AssistantScreenContextController::RequestScreenshot(
                           base::Passed(std::move(layer_owner))));
 }
 
-void AssistantScreenContextController::RequestScreenContext(
-    const gfx::Rect& rect,
-    bool from_user) {
-  // Abort any request in progress and update request state.
-  screen_context_request_factory_.InvalidateWeakPtrs();
-  assistant_screen_context_model_.SetRequestState(
-      ScreenContextRequestState::kInProgress);
-
-  assistant_->RequestScreenContext(
-      rect, from_user,
-      base::BindOnce(
-          &AssistantScreenContextController::OnScreenContextRequestFinished,
-          screen_context_request_factory_.GetWeakPtr()));
-}
-
 void AssistantScreenContextController::OnAssistantControllerConstructed() {
   assistant_controller_->ui_controller()->AddModelObserver(this);
 }
@@ -205,8 +190,8 @@ void AssistantScreenContextController::OnAssistantControllerDestroying() {
 void AssistantScreenContextController::OnUiVisibilityChanged(
     bool visible,
     AssistantSource source) {
-  // We don't initiate a contextual query for caching if the UI is being hidden.
-  // Instead, we abort any requests in progress and reset state.
+  // We don't initiate a cache request if the UI is being hidden. Instead, we
+  // abort any requests in progress and reset state.
   if (!visible) {
     screen_context_request_factory_.InvalidateWeakPtrs();
     assistant_screen_context_model_.SetRequestState(
@@ -218,13 +203,19 @@ void AssistantScreenContextController::OnUiVisibilityChanged(
                                      ->model()
                                      ->input_modality();
 
-  // We don't initiate a contextual query for caching if we are using stylus
-  // input modality because we will do so on metalayer session complete.
+  // We don't initiate a cache request if we are using stylus input modality.
   if (input_modality == InputModality::kStylus)
     return;
 
-  // Request screen context for the entire screen.
-  RequestScreenContext(gfx::Rect(), /*from_user=*/false);
+  // Abort any request in progress and update request state.
+  screen_context_request_factory_.InvalidateWeakPtrs();
+  assistant_screen_context_model_.SetRequestState(
+      ScreenContextRequestState::kInProgress);
+
+  // Cache screen context for the entire screen.
+  assistant_->CacheScreenContext(base::BindOnce(
+      &AssistantScreenContextController::OnScreenContextRequestFinished,
+      screen_context_request_factory_.GetWeakPtr()));
 }
 
 void AssistantScreenContextController::OnScreenContextRequestFinished() {

ash/assistant/assistant_screen_context_controller.h

@@ -44,11 +44,6 @@ class ASH_EXPORT AssistantScreenContextController
   void AddModelObserver(AssistantScreenContextModelObserver* observer);
   void RemoveModelObserver(AssistantScreenContextModelObserver* observer);
 
-  // Requests screen context for the region defined by |rect| (given in DP). If
-  // an empty rect is supplied, the entire screen is captured. If |from_user| is
-  // true, contextual cards will be returned as part of an interaction flow.
-  void RequestScreenContext(const gfx::Rect& rect, bool from_user);
-
   // Requests a screenshot for the region defined by |rect| (given in DP). If
   // an empty rect is supplied, the entire screen is captured. Upon screenshot
   // completion, the specified |callback| is run.

ash/assistant/ui/assistant_container_view_unittest.cc

@@ -67,12 +67,10 @@ class AssistantContainerViewTest : public AshTestBase {
     assistant_binding_->Bind(mojo::MakeRequest(&assistant));
     controller_->SetAssistant(std::move(assistant));
 
-    // Mock any screen context requests by immediately invoking callback.
-    ON_CALL(*assistant_,
-            DoRequestScreenContext(testing::_, testing::_, testing::_))
+    // Mock any screen context cache requests by immediately invoking callback.
+    ON_CALL(*assistant_, DoCacheScreenContext(testing::_))
         .WillByDefault(testing::Invoke(
-            [](const gfx::Rect& rect, bool from_user,
-               base::OnceClosure* callback) { std::move(*callback).Run(); }));
+            [](base::OnceClosure* callback) { std::move(*callback).Run(); }));
   }
 
   base::test::ScopedFeatureList scoped_feature_list_;

chromeos/services/assistant/assistant_manager_service_impl.cc

@@ -192,39 +192,6 @@ void AssistantManagerServiceImpl::SendUpdateSettingsUiRequest(
       });
 }
 
-void AssistantManagerServiceImpl::RequestScreenContext(
-    const gfx::Rect& region,
-    bool from_user,
-    RequestScreenContextCallback callback) {
-  if (!assistant_enabled_ || !context_enabled_) {
-    // If context is not enabled, we notify assistant immediately to activate
-    // the UI.
-    std::move(callback).Run();
-    return;
-  }
-
-  // We wait for the closure to execute twice: once for the screenshot and once
-  // for the view hierarchy. If the screen context request was user initiated,
-  // we then proceed to send a screen context query to the server. Otherwise,
-  // we cache our screen context to send alongside subsequent voice/text
-  // queries. In all cases, |callback| will be run to inform our caller that
-  // context grabbing has been completed.
-  auto on_done = base::BarrierClosure(
-      2, base::BindOnce(
-             from_user ? &AssistantManagerServiceImpl::SendScreenContextQuery
-                       : &AssistantManagerServiceImpl::CacheScreenContext,
-             weak_factory_.GetWeakPtr(), std::move(callback)));
-
-  service_->client()->RequestAssistantStructure(
-      base::BindOnce(&AssistantManagerServiceImpl::OnAssistantStructureReceived,
-                     weak_factory_.GetWeakPtr(), on_done));
-
-  service_->assistant_controller()->RequestScreenshot(
-      region, base::BindOnce(
-                  &AssistantManagerServiceImpl::OnAssistantScreenshotReceived,
-                  weak_factory_.GetWeakPtr(), on_done));
-}
-
 void AssistantManagerServiceImpl::StartVoiceInteraction() {
   platform_api_.SetMicState(true);
   assistant_manager_->StartAssistantInteraction();
@@ -235,6 +202,32 @@ void AssistantManagerServiceImpl::StopActiveInteraction() {
   assistant_manager_->StopAssistantInteraction();
 }
 
+void AssistantManagerServiceImpl::StartCachedScreenContextInteraction() {
+  if (!assistant_enabled_ || !context_enabled_)
+    return;
+
+  // It is illegal to call this method without having first cached screen
+  // context (see CacheScreenContext()).
+  DCHECK(assistant_extra_);
+  DCHECK(assistant_tree_);
+  DCHECK(!assistant_screenshot_.empty());
+
+  SendScreenContextRequest(std::move(assistant_extra_),
+                           std::move(assistant_tree_), assistant_screenshot_);
+}
+
+void AssistantManagerServiceImpl::StartMetalayerInteraction(
+    const gfx::Rect& region) {
+  if (!assistant_enabled_ || !context_enabled_)
+    return;
+
+  service_->assistant_controller()->RequestScreenshot(
+      region,
+      base::BindOnce(&AssistantManagerServiceImpl::SendScreenContextRequest,
+                     weak_factory_.GetWeakPtr(), /*assistant_extra=*/nullptr,
+                     /*assistant_tree=*/nullptr));
+}
+
 void AssistantManagerServiceImpl::SendTextQuery(const std::string& query) {
   assistant_client::VoicelessOptions options;
   options.is_user_initiated = true;
@@ -820,7 +813,33 @@ void AssistantManagerServiceImpl::OnSpeechLevelUpdatedOnMainThread(
       [&speech_level](auto* ptr) { ptr->OnSpeechLevelUpdated(speech_level); });
 }
 
-void AssistantManagerServiceImpl::OnAssistantStructureReceived(
+void AssistantManagerServiceImpl::CacheScreenContext(
+    CacheScreenContextCallback callback) {
+  if (!assistant_enabled_ || !context_enabled_) {
+    std::move(callback).Run();
+    return;
+  }
+
+  // Our callback should be run only after both view hierarchy and screenshot
+  // data have been cached from their respective providers.
+  auto on_done =
+      base::BarrierClosure(2, base::BindOnce(
+                                  [](CacheScreenContextCallback callback) {
+                                    std::move(callback).Run();
+                                  },
+                                  base::Passed(std::move(callback))));
+
+  service_->client()->RequestAssistantStructure(
+      base::BindOnce(&AssistantManagerServiceImpl::CacheAssistantStructure,
+                     weak_factory_.GetWeakPtr(), on_done));
+
+  service_->assistant_controller()->RequestScreenshot(
+      gfx::Rect(),
+      base::BindOnce(&AssistantManagerServiceImpl::CacheAssistantScreenshot,
+                     weak_factory_.GetWeakPtr(), on_done));
+}
+
+void AssistantManagerServiceImpl::CacheAssistantStructure(
     base::OnceClosure on_done,
     ax::mojom::AssistantExtraPtr assistant_extra,
     std::unique_ptr<ui::AssistantTree> assistant_tree) {
@@ -829,41 +848,22 @@ void AssistantManagerServiceImpl::OnAssistantStructureReceived(
   std::move(on_done).Run();
 }
 
-void AssistantManagerServiceImpl::OnAssistantScreenshotReceived(
+void AssistantManagerServiceImpl::CacheAssistantScreenshot(
     base::OnceClosure on_done,
-    const std::vector<uint8_t>& jpg_image) {
-  assistant_screenshot_ = jpg_image;
+    const std::vector<uint8_t>& assistant_screenshot) {
+  assistant_screenshot_ = assistant_screenshot;
   std::move(on_done).Run();
 }
 
-void AssistantManagerServiceImpl::CacheScreenContext(
-    RequestScreenContextCallback callback) {
-  // TODO(dmblack): Cache view hierarchy and screenshot data so that we can
-  // send it alongside subsequent voice/text queries for additional context.
-  assistant_extra_.reset();
-  assistant_tree_.reset();
-  assistant_screenshot_.clear();
-
-  // Run callback now that screen context has been successfully grabbed.
-  std::move(callback).Run();
-}
-
-void AssistantManagerServiceImpl::SendScreenContextQuery(
-    RequestScreenContextCallback callback) {
-  // Note: This call initiates a user facing interaction which will interrupt
-  // any other user interaction in progress. As such, it should only be called
-  // as a direct result of a query for screen related content initiated by the
-  // user, e.g. via suggestion chip press or Assistant stylus tool selection.
+void AssistantManagerServiceImpl::SendScreenContextRequest(
+    ax::mojom::AssistantExtraPtr assistant_extra,
+    std::unique_ptr<ui::AssistantTree> assistant_tree,
+    const std::vector<uint8_t>& assistant_screenshot) {
   assistant_manager_internal_->SendScreenContextRequest(
-      {CreateContextProto(AssistantBundle{
-           std::move(assistant_extra_), std::move(assistant_tree_),
-       }),
-       CreateContextProto(assistant_screenshot_)});
-
+      {CreateContextProto(AssistantBundle{std::move(assistant_extra),
+                                          std::move(assistant_tree)}),
+       CreateContextProto(assistant_screenshot)});
   assistant_screenshot_.clear();
-
-  // Run callback now that screen context has been successfully grabbed.
-  std::move(callback).Run();
 }
 
 }  // namespace assistant

chromeos/services/assistant/assistant_manager_service_impl.h

@@ -80,6 +80,8 @@ class AssistantManagerServiceImpl
       UpdateSettingsUiResponseCallback callback) override;
 
   // mojom::Assistant overrides:
+  void StartCachedScreenContextInteraction() override;
+  void StartMetalayerInteraction(const gfx::Rect& region) override;
   void StartVoiceInteraction() override;
   void StopActiveInteraction() override;
   void SendTextQuery(const std::string& query) override;
@@ -91,9 +93,7 @@ class AssistantManagerServiceImpl
                             int action_index) override;
   void DismissNotification(
       mojom::AssistantNotificationPtr notification) override;
-  void RequestScreenContext(const gfx::Rect& region,
-                            bool from_user,
-                            RequestScreenContextCallback callback) override;
+  void CacheScreenContext(CacheScreenContextCallback callback) override;
 
   // AssistantActionObserver overrides:
   void OnShowHtml(const std::string& html) override;
@@ -135,7 +135,7 @@ class AssistantManagerServiceImpl
       ash::mojom::AssistantAllowedState state) override {}
   void OnLocaleChanged(const std::string& locale) override {}
 
-  // AddDeviceStateListener overrides
+  // assistant_client::DeviceStateListener overrides:
   void OnStartFinished() override;
 
  private:
@@ -176,15 +176,19 @@ class AssistantManagerServiceImpl
 
   void RegisterFallbackMediaHandler();
 
-  void CacheScreenContext(RequestScreenContextCallback callback);
-  void SendScreenContextQuery(RequestScreenContextCallback callback);
-
-  void OnAssistantStructureReceived(
+  void CacheAssistantStructure(
       base::OnceClosure on_done,
       ax::mojom::AssistantExtraPtr assistant_extra,
       std::unique_ptr<ui::AssistantTree> assistant_tree);
-  void OnAssistantScreenshotReceived(base::OnceClosure on_done,
-                                     const std::vector<uint8_t>& jpg_image);
+
+  void CacheAssistantScreenshot(
+      base::OnceClosure on_done,
+      const std::vector<uint8_t>& assistant_screenshot);
+
+  void SendScreenContextRequest(
+      ax::mojom::AssistantExtraPtr assistant_extra,
+      std::unique_ptr<ui::AssistantTree> assistant_tree,
+      const std::vector<uint8_t>& assistant_screenshot);
 
   State state_ = State::STOPPED;
   PlatformApiImpl platform_api_;

chromeos/services/assistant/fake_assistant_manager_service_impl.cc

@@ -46,10 +46,10 @@ void FakeAssistantManagerServiceImpl::SendUpdateSettingsUiRequest(
     const std::string& update,
     UpdateSettingsUiResponseCallback callback) {}
 
-void FakeAssistantManagerServiceImpl::RequestScreenContext(
-    const gfx::Rect& region,
-    bool from_user,
-    RequestScreenContextCallback callback) {}
+void FakeAssistantManagerServiceImpl::StartCachedScreenContextInteraction() {}
+
+void FakeAssistantManagerServiceImpl::StartMetalayerInteraction(
+    const gfx::Rect& region) {}
 
 void FakeAssistantManagerServiceImpl::StartVoiceInteraction() {}
 
@@ -70,5 +70,8 @@ void FakeAssistantManagerServiceImpl::RetrieveNotification(
 void FakeAssistantManagerServiceImpl::DismissNotification(
     mojom::AssistantNotificationPtr notification) {}
 
+void FakeAssistantManagerServiceImpl::CacheScreenContext(
+    CacheScreenContextCallback callback) {}
+
 }  // namespace assistant
 }  // namespace chromeos

chromeos/services/assistant/fake_assistant_manager_service_impl.h

@@ -39,6 +39,8 @@ class FakeAssistantManagerServiceImpl : public AssistantManagerService {
       UpdateSettingsUiResponseCallback callback) override;
 
   // mojom::Assistant overrides:
+  void StartCachedScreenContextInteraction() override;
+  void StartMetalayerInteraction(const gfx::Rect& region) override;
   void StartVoiceInteraction() override;
   void StopActiveInteraction() override;
   void SendTextQuery(const std::string& query) override;
@@ -50,9 +52,7 @@ class FakeAssistantManagerServiceImpl : public AssistantManagerService {
                             int action_index) override;
   void DismissNotification(
       mojom::AssistantNotificationPtr notification) override;
-  void RequestScreenContext(const gfx::Rect& region,
-                            bool from_user,
-                            RequestScreenContextCallback callback) override;
+  void CacheScreenContext(CacheScreenContextCallback callback) override;
 
  private:
   State state_ = State::STOPPED;

chromeos/services/assistant/public/mojom/assistant.mojom

@@ -12,6 +12,17 @@ import "mojo/public/mojom/base/time.mojom";
 
 // Interface to communicate with assistant backend.
 interface Assistant {
+  // Starts a cached screen context interaction. Results related to the screen
+  // context will be returned through the |AssistantInteractionSubscriber|
+  // interface to registered subscribers. It is illegal to call this method
+  // without having first cached screen context (see CacheScreenContext()).
+  StartCachedScreenContextInteraction();
+
+  // Starts a metalayer interaction for the selected screen |region|. Results
+  // related to the selected region will be returned through the
+  // |AssistantInteractionSubscriber| interface to registered subscribers.
+  StartMetalayerInteraction(gfx.mojom.Rect region);
+
   // Starts a new Assistant voice interaction.
   StartVoiceInteraction();
 
@@ -41,15 +52,11 @@ interface Assistant {
   // Dismisses a notification.
   DismissNotification(AssistantNotification notification);
 
-  // Asks assistant to grab screen context. |region| controls the part of the
-  // screen to snapshot. If |region| is empty, it takes fullscreen screenshot.
-  // The |from_user| parameter is true if the request originated from the user,
-  // false otherwise. A user request for screen context may come as a result of
-  // a suggestion chip press or Assistant stylus selection. A request that did
-  // not originate from the user occurs on launch of Assistant to provide
-  // additional context for subsequent voice/text queries. After grab
-  // completion, the callback will be invoked.
-  RequestScreenContext(gfx.mojom.Rect region, bool from_user) => ();
+  // Caches screen context, made up of view hierarchy and screenshot data.
+  // Screen context is used to provide additional context alongside text and
+  // voice queries, and may also be used in standalone screen context
+  // interactions (see StartCachedScreenContextInteraction()).
+  CacheScreenContext() => ();
 };
 
 // Subscribes to assistant's interaction event. These events are server driven

chromeos/services/assistant/test_support/mock_assistant.h

@@ -21,6 +21,10 @@ class MockAssistant : public mojom::Assistant {
   MockAssistant();
   ~MockAssistant() override;
 
+  MOCK_METHOD0(StartCachedScreenContextInteraction, void());
+
+  MOCK_METHOD1(StartMetalayerInteraction, void(const gfx::Rect&));
+
   MOCK_METHOD0(StartVoiceInteraction, void());
 
   MOCK_METHOD0(StopActiveInteraction, void());
@@ -41,17 +45,14 @@ class MockAssistant : public mojom::Assistant {
   MOCK_METHOD1(DismissNotification,
                void(chromeos::assistant::mojom::AssistantNotificationPtr));
 
-  // Mock DoRequestScreenContext in lieu of RequestScreenContext.
-  MOCK_METHOD3(DoRequestScreenContext,
-               void(const gfx::Rect&, bool from_user, base::OnceClosure*));
+  // Mock DoCacheScreenContext in lieu of CacheScreenContext.
+  MOCK_METHOD1(DoCacheScreenContext, void(base::OnceClosure*));
 
-  // Note: We can't mock RequestScreenContext directly because of the move
+  // Note: We can't mock CacheScreenContext directly because of the move
   // semantics required around base::OnceClosure. Instead, we route calls to a
-  // mockable delegate method, DoRequestScreenContext.
-  void RequestScreenContext(const gfx::Rect& region,
-                            bool from_user,
-                            base::OnceClosure callback) override {
-    DoRequestScreenContext(region, from_user, &callback);
+  // mockable delegate method, DoCacheScreenContext.
+  void CacheScreenContext(base::OnceClosure callback) override {
+    DoCacheScreenContext(&callback);
   }
 
  private: