[Contextual Tasks] Upload tab context before sending Lens side panel queries

Refactors `LensQueryFlowRouter::SendContextualTextQuery` to create a
`CreateSearchUrlRequestInfo` with type `kAim` and route it through
`SendInteractionToContextualTasks`.

Modifies `SendInteractionToContextualTasks` to ensure a
`ContextualSearchSessionHandle` exists and has tab context before
generating the search URL. If a session is missing, it is created, and
the search request is cached while `TabContextualizationController`
fetches and uploads the page context.

Bug: 467770833
Change-Id: I4e0d2ccd224fcdc9878d5afdb622e3dc340d3ea8
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7255015
Reviewed-by: Duncan Mercer <mercerd@google.com>
Commit-Queue: Juan Mojica <juanmojica@google.com>
Cr-Commit-Position: refs/heads/main@{#1558271}
diff --git a/chrome/browser/contextual_tasks/contextual_tasks_ui_service.cc b/chrome/browser/contextual_tasks/contextual_tasks_ui_service.cc
index f6a1b4d..4f536d0 100644
--- a/chrome/browser/contextual_tasks/contextual_tasks_ui_service.cc
+++ b/chrome/browser/contextual_tasks/contextual_tasks_ui_service.cc
@@ -579,7 +579,7 @@
     // Associate the web contents with the task and set the session handle if
     // provided.
     content::WebContents* web_contents = coordinator->GetActiveWebContents();
-    AssociateWebContentsToTask(panel_contents, task.GetTaskId());
+    AssociateWebContentsToTask(web_contents, task.GetTaskId());
     if (session_handle) {
       ContextualSearchWebContentsHelper::GetOrCreateForWebContents(web_contents)
           ->set_session_handle(std::move(session_handle));
diff --git a/chrome/browser/ui/lens/lens_query_flow_router.cc b/chrome/browser/ui/lens/lens_query_flow_router.cc
index 08d652b..139369f 100644
--- a/chrome/browser/ui/lens/lens_query_flow_router.cc
+++ b/chrome/browser/ui/lens/lens_query_flow_router.cc
@@ -22,6 +22,7 @@
 #include "components/lens/lens_overlay_mime_type.h"
 #include "components/lens/lens_url_utils.h"
 #include "components/lens/ref_counted_lens_overlay_client_logs.h"
+#include "components/lens/tab_contextualization_controller.h"
 #include "components/sessions/content/session_tab_helper.h"
 #include "net/base/url_util.h"
 
@@ -131,7 +132,11 @@
     lens::LensOverlaySelectionType lens_selection_type,
     std::map<std::string, std::string> additional_search_query_params) {
   if (contextual_tasks::GetEnableLensInContextualTasks()) {
-    LoadQueryInContextualTasks(query_text);
+    auto request_info = CreateSearchUrlRequestInfoFromInteraction(
+        /*region=*/nullptr, /*region_bytes=*/std::nullopt, query_text,
+        lens_selection_type, additional_search_query_params, query_start_time);
+    request_info->search_url_type = SearchUrlType::kAim;
+    SendInteractionToContextualTasks(std::move(request_info));
     return;
   }
 
@@ -176,20 +181,24 @@
       ->initial_screenshot();
 }
 
-void LensQueryFlowRouter::LoadQueryInContextualTasks(
-    const std::string& query_text) {
-  auto* ui_service =
-      contextual_tasks::ContextualTasksUiServiceFactory::GetForBrowserContext(
-          web_contents()->GetBrowserContext());
-  GURL ai_url = ui_service->GetDefaultAiPageUrl();
-  ai_url = net::AppendQueryParameter(ai_url, "q", query_text);
-  OpenContextualTasksPanel(ai_url);
-}
-
 void LensQueryFlowRouter::SendInteractionToContextualTasks(
     std::unique_ptr<CreateSearchUrlRequestInfo> request_info) {
-  auto* session_handle = GetOrCreateContextualSearchSessionHandle();
-  session_handle->CreateSearchUrl(
+  // If there is no existing session handle, then the search URL request will
+  // need to wait for the tab context to be added before being sent.
+  if (!GetContextualSearchSessionHandle()) {
+    pending_session_handle_ = CreateContextualSearchSessionHandle();
+    pending_session_handle_->NotifySessionStarted();
+    pending_search_url_request_ = std::move(request_info);
+    // Upload the page context when creating a session handle.
+    if (auto* controller =
+            TabContextualizationController::From(tab_interface())) {
+      controller->GetPageContext(base::BindOnce(
+          &LensQueryFlowRouter::UploadContextualInputData,
+          weak_factory_.GetWeakPtr(), pending_session_handle_.get()));
+    }
+    return;
+  }
+  GetContextualSearchSessionHandle()->CreateSearchUrl(
       std::move(request_info),
       base::BindOnce(&LensQueryFlowRouter::OpenContextualTasksPanel,
                      weak_factory_.GetWeakPtr()));
@@ -231,6 +240,17 @@
 
   session_handle->StartTabContextUploadFlow(
       token, std::move(contextual_input_data), std::move(image_options));
+
+  // While a pending search URL request does not need to wait for the tab
+  // context to upload, it does need to wait for tab context to be added to the
+  // session handle before creating the search URL so it is properly
+  // contextualized.
+  if (pending_search_url_request_) {
+    session_handle->CreateSearchUrl(
+        std::move(pending_search_url_request_),
+        base::BindOnce(&LensQueryFlowRouter::OpenContextualTasksPanel,
+                       weak_factory_.GetWeakPtr()));
+  }
 }
 
 std::unique_ptr<lens::ContextualInputData>
@@ -291,18 +311,6 @@
 }
 
 contextual_search::ContextualSearchSessionHandle*
-LensQueryFlowRouter::GetOrCreateContextualSearchSessionHandle() {
-  auto* session_handle = GetContextualSearchSessionHandle();
-  if (session_handle) {
-    return session_handle;
-  }
-
-  pending_session_handle_ = CreateContextualSearchSessionHandle();
-  pending_session_handle_->NotifySessionStarted();
-  return pending_session_handle_.get();
-}
-
-contextual_search::ContextualSearchSessionHandle*
 LensQueryFlowRouter::GetContextualSearchSessionHandle() const {
   if (pending_session_handle_) {
     return pending_session_handle_.get();
@@ -320,8 +328,8 @@
     return nullptr;
   }
 
-  auto* helper = ContextualSearchWebContentsHelper::FromWebContents(
-      coordinator->GetActiveWebContents());
+  auto* helper =
+      ContextualSearchWebContentsHelper::FromWebContents(web_contents);
   if (helper && helper->session_handle()) {
     return helper->session_handle();
   }
diff --git a/chrome/browser/ui/lens/lens_query_flow_router.h b/chrome/browser/ui/lens/lens_query_flow_router.h
index 4a878de..3e018c3 100644
--- a/chrome/browser/ui/lens/lens_query_flow_router.h
+++ b/chrome/browser/ui/lens/lens_query_flow_router.h
@@ -119,9 +119,6 @@
     return Profile::FromBrowserContext(web_contents()->GetBrowserContext());
   }
 
-  // Loads the provided query text in the contextual tasks panel.
-  void LoadQueryInContextualTasks(const std::string& query_text);
-
   // Sends the provided request info to the contextual tasks panel to create a
   // search URL which is then loaded into the contextual tasks panel.
   void SendInteractionToContextualTasks(
@@ -165,16 +162,15 @@
       std::map<std::string, std::string> additional_search_query_params,
       base::Time query_start_time);
 
-  // Returns the contextual search session handle for the query router. If the
-  // handle does not exist, it will create one.
-  contextual_search::ContextualSearchSessionHandle*
-  GetOrCreateContextualSearchSessionHandle();
-
   // Returns the contextual search session handle for the query router if it
   // exists.
   contextual_search::ContextualSearchSessionHandle*
   GetContextualSearchSessionHandle() const;
 
+  // Stores a pending search request to be sent to contextual tasks after the
+  // tab context is ready.
+  std::unique_ptr<CreateSearchUrlRequestInfo> pending_search_url_request_;
+
   // The contextual search session handle that is used to make requests to the
   // contextual search service. This is only stored by this query router in
   // cases where the overlay has been opened but a results panel is not present.
diff --git a/chrome/browser/ui/lens/lens_query_flow_router_unittest.cc b/chrome/browser/ui/lens/lens_query_flow_router_unittest.cc
index e5c2db5..0293c91b 100644
--- a/chrome/browser/ui/lens/lens_query_flow_router_unittest.cc
+++ b/chrome/browser/ui/lens/lens_query_flow_router_unittest.cc
@@ -21,6 +21,7 @@
 #include "components/lens/contextual_input.h"
 #include "components/lens/lens_features.h"
 #include "components/lens/lens_url_utils.h"
+#include "components/lens/tab_contextualization_controller.h"
 #include "components/tabs/public/mock_tab_interface.h"
 #include "content/public/test/browser_task_environment.h"
 #include "content/public/test/test_renderer_host.h"
@@ -129,6 +130,19 @@
       pending_mock_session_handle_;
 };
 
+class MockTabContextualizationController
+    : public TabContextualizationController {
+ public:
+  explicit MockTabContextualizationController(tabs::TabInterface* tab)
+      : TabContextualizationController(tab) {}
+  ~MockTabContextualizationController() override = default;
+
+  MOCK_METHOD(void,
+              GetPageContext,
+              (GetPageContextCallback callback),
+              (override));
+};
+
 class MockContextualTasksUiService
     : public contextual_tasks::ContextualTasksUiService {
  public:
@@ -373,7 +387,19 @@
         ->SetTestingFactory(
             profile_.get(),
             base::BindRepeating(&CreateMockContextualTasksUiService));
+    mock_tab_contextualization_controller_ =
+        std::make_unique<MockTabContextualizationController>(
+            &mock_tab_interface_);
   }
+
+  void TearDown() override {
+    // Controller must be destroyed before the tab interface and user data host.
+    mock_tab_contextualization_controller_.reset();
+    LensQueryFlowRouterTest::TearDown();
+  }
+
+  std::unique_ptr<MockTabContextualizationController>
+      mock_tab_contextualization_controller_;
 };
 
 TEST_F(LensQueryFlowRouterContextualTaskEnabledTest,
@@ -450,6 +476,9 @@
   // Assert: Create expectation to call CreateSearchUrl. We also expect a call
   // to open the side panel, but that is harder to mock, so we omit it for now.
   EXPECT_CALL(*router.mock_session_handle(), NotifySessionStarted());
+  // StartTabContextUploadFlow is called as part of OnFinishedAddingTabContext.
+  EXPECT_CALL(*router.mock_session_handle(),
+              StartTabContextUploadFlow(_, _, _));
   EXPECT_CALL(
       *router.mock_session_handle(),
       CreateSearchUrl(
@@ -465,6 +494,9 @@
                   GURL("https://www.google.com/search?q=test"),
                   testing::Pointer(router.mock_session_handle())))
       .Times(1);
+  EXPECT_CALL(*mock_tab_contextualization_controller_, GetPageContext(_))
+      .WillOnce([](lens::TabContextualizationController::GetPageContextCallback
+                       callback) { std::move(callback).Run(nullptr); });
 
   // Act: Call the method.
   router.SendRegionSearch(query_start_time, std::move(region), selection_type,
@@ -498,6 +530,9 @@
 
   // Assert: Create expectation to call CreateSearchUrl.
   EXPECT_CALL(*router.mock_session_handle(), NotifySessionStarted());
+  // StartTabContextUploadFlow is called as part of OnFinishedAddingTabContext.
+  EXPECT_CALL(*router.mock_session_handle(),
+              StartTabContextUploadFlow(_, _, _));
   EXPECT_CALL(
       *router.mock_session_handle(),
       CreateSearchUrl(
@@ -513,6 +548,9 @@
                   GURL("https://www.google.com/search?q=test"),
                   testing::Pointer(router.mock_session_handle())))
       .Times(1);
+  EXPECT_CALL(*mock_tab_contextualization_controller_, GetPageContext(_))
+      .WillOnce([](lens::TabContextualizationController::GetPageContextCallback
+                       callback) { std::move(callback).Run(nullptr); });
 
   // Act: Call the method.
   router.SendTextOnlyQuery(query_start_time, query_text, selection_type,
@@ -531,17 +569,39 @@
       lens::LensOverlaySelectionType::MULTIMODAL_SUGGEST_TYPEAHEAD;
   std::map<std::string, std::string> additional_params;
 
-  // Assert: Create expectation to call GetDefaultAiPageUrl.
+  // Arrange: Create expected request info.
+  auto expected_request_info = std::make_unique<CreateSearchUrlRequestInfo>();
+  expected_request_info->search_url_type =
+      contextual_search::ContextualSearchContextController::SearchUrlType::kAim;
+  expected_request_info->query_text = query_text;
+  expected_request_info->query_start_time = query_start_time;
+  expected_request_info->lens_overlay_selection_type = selection_type;
+  expected_request_info->additional_params = additional_params;
+  expected_request_info->image_crop = std::nullopt;
+
+  // Assert: Create expectation to call CreateSearchUrl.
+  EXPECT_CALL(*router.mock_session_handle(), NotifySessionStarted());
+  // StartTabContextUploadFlow is called as part of OnFinishedAddingTabContext.
+  EXPECT_CALL(*router.mock_session_handle(),
+              StartTabContextUploadFlow(_, _, _));
+  EXPECT_CALL(
+      *router.mock_session_handle(),
+      CreateSearchUrl(
+          CreateSearchUrlRequestInfoMatches(expected_request_info.get()), _))
+      .WillOnce(base::test::RunOnceCallback<1>(
+          GURL("https://www.google.com/search?q=test")));
   auto* service = static_cast<MockContextualTasksUiService*>(
       contextual_tasks::ContextualTasksUiServiceFactory::GetForBrowserContext(
           profile_.get()));
-  EXPECT_CALL(*service, GetDefaultAiPageUrl())
-      .WillOnce(Return(GURL("https://example.com")));
   EXPECT_CALL(*service,
               StartTaskUiInSidePanel(
                   mock_browser_window_interface_.get(), &mock_tab_interface_,
-                  GURL("https://example.com/?q=test+query"), testing::IsNull()))
+                  GURL("https://www.google.com/search?q=test"),
+                  testing::Pointer(router.mock_session_handle())))
       .Times(1);
+  EXPECT_CALL(*mock_tab_contextualization_controller_, GetPageContext(_))
+      .WillOnce([](lens::TabContextualizationController::GetPageContextCallback
+                       callback) { std::move(callback).Run(nullptr); });
 
   // Act: Call the method.
   router.SendContextualTextQuery(query_start_time, query_text, selection_type,
@@ -576,6 +636,9 @@
   // Assert: Create expectation to call CreateSearchUrl. We also expect a call
   // to open the side panel, but that is harder to mock, so we omit it for now.
   EXPECT_CALL(*router.mock_session_handle(), NotifySessionStarted());
+  // StartTabContextUploadFlow is called as part of OnFinishedAddingTabContext.
+  EXPECT_CALL(*router.mock_session_handle(),
+              StartTabContextUploadFlow(_, _, _));
   EXPECT_CALL(
       *router.mock_session_handle(),
       CreateSearchUrl(
@@ -591,6 +654,9 @@
                   GURL("https://www.google.com/search?q=test"),
                   testing::Pointer(router.mock_session_handle())))
       .Times(1);
+  EXPECT_CALL(*mock_tab_contextualization_controller_, GetPageContext(_))
+      .WillOnce([](lens::TabContextualizationController::GetPageContextCallback
+                       callback) { std::move(callback).Run(nullptr); });
 
   // Act: Call the method.
   router.SendMultimodalRequest(query_start_time, std::move(region), query_text,