MultiTabContent conversations can each be associated with different s…

…ets of content

browser/brave_tab_helpers.cc

@@ -140,6 +140,8 @@ void AttachTabHelpers(content::WebContents* web_contents) {
   if (ai_chat::IsAllowedForContext(context)) {
     ai_chat::AIChatTabHelper::CreateForWebContents(
         web_contents,
+        ai_chat::AIChatServiceFactory::GetForBrowserContext(
+            web_contents->GetBrowserContext()),
 #if BUILDFLAG(ENABLE_PRINT_PREVIEW)
         std::make_unique<ai_chat::PrintPreviewExtractor>(web_contents)
 #else

browser/ui/BUILD.gn

@@ -87,8 +87,6 @@ source_set("ui") {
 
   if (enable_ai_chat) {
     sources += [
-      "ai_chat/associated_multi_tab_content.cc",
-      "ai_chat/associated_multi_tab_content.h",
       "webui/ai_chat/ai_chat_ui.cc",
       "webui/ai_chat/ai_chat_ui.h",
       "webui/ai_chat/ai_chat_ui_page_handler.cc",

browser/ui/webui/ai_chat/ai_chat_ui_page_handler.cc

@@ -9,13 +9,15 @@
 #include <utility>
 #include <vector>
 
+#include "base/containers/fixed_flat_set.h"
 #include "base/memory/weak_ptr.h"
 #include "base/strings/utf_string_conversions.h"
 #include "brave/browser/ai_chat/ai_chat_service_factory.h"
 #include "brave/browser/ui/side_panel/ai_chat/ai_chat_side_panel_utils.h"
 #include "brave/components/ai_chat/content/browser/ai_chat_tab_helper.h"
 #include "brave/components/ai_chat/core/browser/ai_chat_service.h"
 #include "brave/components/ai_chat/core/browser/associated_content_driver.h"
+#include "brave/components/ai_chat/core/browser/associated_multi_tab_content.h"
 #include "brave/components/ai_chat/core/browser/constants.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h"
@@ -48,6 +50,8 @@ constexpr char kURLGoPremium[] =
     "https://account.brave.com/account/?intent=checkout&product=leo";
 constexpr char kURLManagePremium[] = "https://account.brave.com/";
 #endif
+static const auto kAllowedSchemes = base::MakeFixedFlatSet<std::string_view>(
+    {url::kHttpsScheme, url::kHttpScheme, url::kFileScheme, url::kDataScheme});
 }  // namespace
 
 namespace ai_chat {
@@ -196,22 +200,7 @@ void AIChatUIPageHandler::SetChatUI(
     mojo::PendingRemote<mojom::ChatUI> chat_ui) {
   chat_ui_.Bind(std::move(chat_ui));
 
-  std::vector<mojom::WebSiteInfoDetailPtr> tabs;
-  if (auto* browser = chrome::FindBrowserWithTab(owner_web_contents_)) {
-    for (int i = 0; i < browser->tab_strip_model()->count(); ++i) {
-      auto* wc = browser->tab_strip_model()->GetWebContentsAt(i);
-      if (wc == owner_web_contents_) {
-        continue;
-      }
-      auto info = mojom::WebSiteInfoDetail::New();
-      info->hostname = wc->GetVisibleURL().host();
-      info->title = base::UTF16ToUTF8(wc->GetTitle());
-      info->url = wc->GetURL();
-
-      tabs.push_back(std::move(info));
-    }
-  }
-  chat_ui_->SetInitialData(active_chat_tab_helper_ == nullptr, std::move(tabs));
+  chat_ui_->SetInitialData(active_chat_tab_helper_ == nullptr);
 }
 
 void AIChatUIPageHandler::BindRelatedConversation(
@@ -254,20 +243,41 @@ void AIChatUIPageHandler::NewConversation(
 void AIChatUIPageHandler::NewMultiTabConversation(
     mojo::PendingReceiver<mojom::ConversationHandler> receiver,
     mojo::PendingRemote<mojom::ConversationUI> conversation_ui_handler) {
-  ConversationHandler* conversation;
   // Make an associated content for all tabs in current window
   Browser* browser = ai_chat::GetBrowserForWebContents(owner_web_contents_);
   if (!browser) {
     return;
   }
 
-  auto* multi_tab_content = browser->browser_window_features()
-                                ->ai_chat_associated_multi_tab_content();
+  ConversationHandler* conversation =
+      AIChatServiceFactory::GetForBrowserContext(profile_)
+          ->CreateConversation();
+
+  TabStripModel* tab_strip_model = browser->tab_strip_model();
+  std::vector<AssociatedContentDriver*> starting_web_contentses;
+
+  for (int i = 0; i < tab_strip_model->count(); ++i) {
+    content::WebContents* web_contents = tab_strip_model->GetWebContentsAt(i);
+    // Temporary fix to ensure that the tab is loaded before we try to get the
+    // content
+    web_contents->GetController().LoadIfNecessary();
+    AIChatTabHelper* associated_content =
+        ai_chat::AIChatTabHelper::FromWebContents(web_contents);
+    if (base::Contains(kAllowedSchemes,
+                       associated_content->GetURL().scheme())) {
+      DVLOG(2) << "Tab " << i
+               << " URL: " << associated_content->GetURL().spec();
+      starting_web_contentses.push_back(associated_content);
+    }
+  }
 
-  conversation = AIChatServiceFactory::GetForBrowserContext(profile_)
-                     ->CreateConversation();
+  auto multi_tab_content = std::make_unique<AssociatedMultiTabContent>(
+      std::move(starting_web_contentses),
+      owner_web_contents_->GetBrowserContext()
+          ->GetDefaultStoragePartition()
+          ->GetURLLoaderFactoryForBrowserProcess());
 
-  conversation->SetAssociatedContentDelegate(multi_tab_content->GetWeakPtr());
+  conversation->SetMultiTabContent(std::move(multi_tab_content));
 
   conversation->Bind(std::move(receiver), std::move(conversation_ui_handler));
 }

components/ai_chat/content/browser/ai_chat_tab_helper.cc

@@ -17,6 +17,7 @@
 #include "base/strings/string_util.h"
 #include "brave/components/ai_chat/content/browser/page_content_fetcher.h"
 #include "brave/components/ai_chat/content/browser/pdf_utils.h"
+#include "brave/components/ai_chat/core/browser/ai_chat_service.h"
 #include "brave/components/ai_chat/core/browser/associated_content_driver.h"
 #include "brave/components/ai_chat/core/browser/constants.h"
 #include "brave/components/ai_chat/core/browser/utils.h"
@@ -90,17 +91,20 @@ void AIChatTabHelper::BindPageContentExtractorHost(
 }
 
 AIChatTabHelper::AIChatTabHelper(content::WebContents* web_contents,
+                                 AIChatService* ai_chat_service,
                                  std::unique_ptr<PrintPreviewExtractionDelegate>
                                      print_preview_extraction_delegate)
     : content::WebContentsObserver(web_contents),
       content::WebContentsUserData<AIChatTabHelper>(*web_contents),
       AssociatedContentDriver(web_contents->GetBrowserContext()
                                   ->GetDefaultStoragePartition()
                                   ->GetURLLoaderFactoryForBrowserProcess()),
+      ai_chat_service_(ai_chat_service),
       print_preview_extraction_delegate_(
           std::move(print_preview_extraction_delegate)),
       page_content_fetcher_delegate_(
           std::make_unique<PageContentFetcher>(web_contents)) {
+  ai_chat_service->RegisterAssociatedContentsAvailable(this);
   favicon::ContentFaviconDriver::FromWebContents(web_contents)
       ->AddObserver(this);
   previous_page_title_ = web_contents->GetTitle();
@@ -132,6 +136,7 @@ void AIChatTabHelper::OnPDFA11yInfoLoaded() {
 void AIChatTabHelper::WebContentsDestroyed() {
   favicon::ContentFaviconDriver::FromWebContents(web_contents())
       ->RemoveObserver(this);
+  ai_chat_service_->AssociatedContentsDestroyed(this);
   inner_web_contents_ = nullptr;
 }
 
@@ -353,6 +358,7 @@ std::u16string AIChatTabHelper::GetPageTitle() const {
 
 void AIChatTabHelper::OnNewPage(int64_t navigation_id) {
   DVLOG(3) << __func__ << " id: " << navigation_id;
+  ai_chat_service_->OnContentMetadataChanged();
   AssociatedContentDriver::OnNewPage(navigation_id);
   if (pending_get_page_content_callback_) {
     std::move(pending_get_page_content_callback_).Run("", false, "");

components/ai_chat/content/browser/ai_chat_tab_helper.h

@@ -118,6 +118,7 @@ class AIChatTabHelper : public content::WebContentsObserver,
   // PrintPreviewExtractionDelegate is provided as it's implementation is
   // in a different layer.
   AIChatTabHelper(content::WebContents* web_contents,
+                  AIChatService* ai_chat_service,
                   std::unique_ptr<PrintPreviewExtractionDelegate>
                       print_preview_extraction_delegate);
 
@@ -175,6 +176,7 @@ class AIChatTabHelper : public content::WebContentsObserver,
   void SetPendingGetContentCallback(GetPageContentCallback callback);
 
   raw_ptr<AIChatMetrics> ai_chat_metrics_;
+  raw_ptr<AIChatService> ai_chat_service_;
 
   bool is_same_document_navigation_ = false;
   int pending_navigation_id_;

components/ai_chat/core/browser/BUILD.gn

@@ -23,6 +23,8 @@ static_library("browser") {
     "associated_archive_content.h",
     "associated_content_driver.cc",
     "associated_content_driver.h",
+    "associated_multi_tab_content.cc",
+    "associated_multi_tab_content.h",
     "constants.cc",
     "constants.h",
     "conversation_handler.cc",