[AI Chat]: Manage associations from PageHandler

browser/ai_chat/BUILD.gn

@@ -7,8 +7,6 @@ import("//printing/buildflags/buildflags.gni")
 
 static_library("ai_chat") {
   sources = [
-    "ai_chat_associated_tab_provider.cc",
-    "ai_chat_associated_tab_provider.h",
     "ai_chat_service_factory.cc",
     "ai_chat_service_factory.h",
     "ai_chat_settings_helper.cc",

browser/ai_chat/ai_chat_service_factory.cc

@@ -10,7 +10,6 @@
 #include <utility>
 
 #include "base/no_destructor.h"
-#include "brave/browser/ai_chat/ai_chat_associated_tab_provider.h"
 #include "brave/browser/ai_chat/ai_chat_utils.h"
 #include "brave/browser/brave_browser_process.h"
 #include "brave/browser/misc_metrics/process_misc_metrics.h"
@@ -75,9 +74,7 @@ AIChatServiceFactory::BuildServiceInstanceForBrowserContext(
   return std::make_unique<AIChatService>(
       ModelServiceFactory::GetForBrowserContext(context),
 
-      std::move(credential_manager),
-      std::make_unique<AIChatAssociatedTabProvider>(),
-      user_prefs::UserPrefs::Get(context),
+      std::move(credential_manager), user_prefs::UserPrefs::Get(context),
       (g_brave_browser_process->process_misc_metrics())
           ? g_brave_browser_process->process_misc_metrics()->ai_chat_metrics()
           : nullptr,

browser/ui/webui/ai_chat/ai_chat_ui_page_handler.cc

@@ -99,6 +99,44 @@ void AIChatUIPageHandler::HandleVoiceRecognition(
 #endif
 }
 
+void AIChatUIPageHandler::AssociateTab(mojom::AvailableTabPtr tab,
+                                       const std::string& conversation_uuid) {
+  auto* content = content::WebContents::FromFrameTreeNodeId(
+      static_cast<content::FrameTreeNodeId>(tab->frame_tree_node_id));
+  if (!content) {
+    return;
+  }
+
+  // Ensure the content is loaded before associating
+  content->GetController().LoadIfNecessary();
+
+  auto* tab_helper = ai_chat::AIChatTabHelper::FromWebContents(content);
+  if (!tab_helper) {
+    return;
+  }
+
+  AIChatServiceFactory::GetForBrowserContext(profile_)->AssociateContent(
+      tab_helper, conversation_uuid);
+}
+
+void AIChatUIPageHandler::DisassociateTab(
+    mojom::AvailableTabPtr tab,
+    const std::string& conversation_uuid) {
+  auto* content = content::WebContents::FromFrameTreeNodeId(
+      static_cast<content::FrameTreeNodeId>(tab->frame_tree_node_id));
+  if (!content) {
+    return;
+  }
+
+  auto* tab_helper = ai_chat::AIChatTabHelper::FromWebContents(content);
+  if (!tab_helper) {
+    return;
+  }
+
+  AIChatServiceFactory::GetForBrowserContext(profile_)->DisassociateContent(
+      tab_helper, conversation_uuid);
+}
+
 void AIChatUIPageHandler::OpenAIChatSettings() {
   content::WebContents* contents_to_navigate =
       (active_chat_tab_helper_) ? active_chat_tab_helper_->web_contents()

browser/ui/webui/ai_chat/ai_chat_ui_page_handler.h

@@ -54,6 +54,10 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
   void RefreshPremiumSession() override;
   void ManagePremium() override;
   void HandleVoiceRecognition(const std::string& conversation_uuid) override;
+  void AssociateTab(mojom::AvailableTabPtr tab,
+                    const std::string& conversation_uuid) override;
+  void DisassociateTab(mojom::AvailableTabPtr tab,
+                       const std::string& conversation_uuid) override;
   void CloseUI() override;
   void SetChatUI(mojo::PendingRemote<mojom::ChatUI> chat_ui,
                  SetChatUICallback callback) override;

components/ai_chat/core/browser/BUILD.gn

@@ -22,7 +22,6 @@ static_library("browser") {
     "associated_archive_content.h",
     "associated_content_driver.cc",
     "associated_content_driver.h",
-    "associated_tab_delegate.h",
     "constants.cc",
     "constants.h",
     "conversation_handler.cc",

components/ai_chat/core/browser/ai_chat_service.cc

@@ -93,7 +93,6 @@ bool IsConversationUpdatedTimeWithinRange(
 AIChatService::AIChatService(
     ModelService* model_service,
     std::unique_ptr<AIChatCredentialManager> ai_chat_credential_manager,
-    std::unique_ptr<AssociatedTabDelegate> associated_tab_delegate,
     PrefService* profile_prefs,
     AIChatMetrics* ai_chat_metrics,
     os_crypt_async::OSCryptAsync* os_crypt_async,
@@ -109,7 +108,6 @@ AIChatService::AIChatService(
           std::make_unique<AIChatFeedbackAPI>(url_loader_factory_,
                                               std::string(channel_string))),
       credential_manager_(std::move(ai_chat_credential_manager)),
-      associated_tab_delegate_(std::move(associated_tab_delegate)),
       profile_path_(profile_path) {
   DCHECK(profile_prefs_);
   pref_change_registrar_.Init(profile_prefs_);
@@ -297,6 +295,18 @@ ConversationHandler* AIChatService::CreateConversationHandlerForContent(
   return conversation;
 }
 
+void AIChatService::AssociateContent(AssociatedContentDriver* driver,
+                                     const std::string& conversation_uuid) {
+  DCHECK(base::Contains(conversation_handlers_, conversation_uuid));
+  conversation_handlers_.at(conversation_uuid)->AddAssociation(driver);
+}
+
+void AIChatService::DisassociateContent(AssociatedContentDriver* driver,
+                                        const std::string& conversation_uuid) {
+  DCHECK(base::Contains(conversation_handlers_, conversation_uuid));
+  conversation_handlers_.at(conversation_uuid)->RemoveAssociation(driver);
+}
+
 void AIChatService::DeleteConversations(std::optional<base::Time> begin_time,
                                         std::optional<base::Time> end_time) {
   if (!begin_time.has_value() && !end_time.has_value()) {
@@ -951,13 +961,4 @@ void AIChatService::OpenConversationWithStagedEntries(
   conversation->MaybeFetchOrClearContentStagedConversation();
 }
 
-AssociatedContentDriver* AIChatService::GetAssociatedContent(
-    const mojom::AvailableTabPtr& tab) {
-  // Can be null in tests
-  if (!associated_tab_delegate_) {
-    return nullptr;
-  }
-  return associated_tab_delegate_->GetAssociatedContent(tab);
-}
-
 }  // namespace ai_chat

components/ai_chat/core/browser/ai_chat_service.h

@@ -29,7 +29,6 @@
 #include "brave/components/ai_chat/core/browser/ai_chat_database.h"
 #include "brave/components/ai_chat/core/browser/ai_chat_feedback_api.h"
 #include "brave/components/ai_chat/core/browser/ai_chat_metrics.h"
-#include "brave/components/ai_chat/core/browser/associated_tab_delegate.h"
 #include "brave/components/ai_chat/core/browser/conversation_handler.h"
 #include "brave/components/ai_chat/core/browser/engine/engine_consumer.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h"
@@ -56,7 +55,6 @@ namespace ai_chat {
 
 class ModelService;
 class AIChatMetrics;
-class AvailableTabDelegate;
 class AssociatedContentDriver;
 
 // Main entry point for creating and consuming AI Chat conversations
@@ -70,7 +68,6 @@ class AIChatService : public KeyedService,
   AIChatService(
       ModelService* model_service,
       std::unique_ptr<AIChatCredentialManager> ai_chat_credential_manager,
-      std::unique_ptr<AssociatedTabDelegate> associated_tab_delegate,
       PrefService* profile_prefs,
       AIChatMetrics* ai_chat_metrics,
       os_crypt_async::OSCryptAsync* os_crypt_async,
@@ -126,6 +123,11 @@ class AIChatService : public KeyedService,
       base::WeakPtr<ConversationHandler::AssociatedContentDelegate>
           associated_content);
 
+  void AssociateContent(AssociatedContentDriver* associated_content,
+                        const std::string& conversation_uuid);
+  void DisassociateContent(AssociatedContentDriver* associated_content,
+                           const std::string& conversation_uuid);
+
   // Removes all in-memory and persisted data for all conversations
   void DeleteConversations(std::optional<base::Time> begin_time = std::nullopt,
                            std::optional<base::Time> end_time = std::nullopt);
@@ -141,9 +143,6 @@ class AIChatService : public KeyedService,
           associated_content,
       base::OnceClosure open_ai_chat);
 
-  AssociatedContentDriver* GetAssociatedContent(
-      const mojom::AvailableTabPtr& tab);
-
   // mojom::Service
   void MarkAgreementAccepted() override;
   void EnableStoragePref() override;
@@ -237,7 +236,6 @@ class AIChatService : public KeyedService,
 
   std::unique_ptr<AIChatFeedbackAPI> feedback_api_;
   std::unique_ptr<AIChatCredentialManager> credential_manager_;
-  std::unique_ptr<AssociatedTabDelegate> associated_tab_delegate_;
 
   base::FilePath profile_path_;
 

components/ai_chat/core/browser/conversation_handler.cc

@@ -45,6 +45,7 @@
 #include "brave/components/ai_chat/core/browser/ai_chat_feedback_api.h"
 #include "brave/components/ai_chat/core/browser/ai_chat_service.h"
 #include "brave/components/ai_chat/core/browser/associated_archive_content.h"
+#include "brave/components/ai_chat/core/browser/associated_content_driver.h"
 #include "brave/components/ai_chat/core/browser/local_models_updater.h"
 #include "brave/components/ai_chat/core/browser/model_service.h"
 #include "brave/components/ai_chat/core/browser/model_validator.h"
@@ -541,13 +542,12 @@ void ConversationHandler::GetState(GetStateCallback callback) {
   std::move(callback).Run(std::move(state));
 }
 
-void ConversationHandler::AddAssociatedTab(mojom::AvailableTabPtr tab) {
-  if (HasAnyHistory()) {
+void ConversationHandler::AddAssociation(AssociatedContentDriver* delegate) {
+  if (!delegate) {
     return;
   }
 
-  auto* tab_delegate = ai_chat_service_->GetAssociatedContent(tab);
-  if (!tab_delegate) {
+  if (HasAnyHistory()) {
     return;
   }
 
@@ -557,20 +557,20 @@ void ConversationHandler::AddAssociatedTab(mojom::AvailableTabPtr tab) {
   }
 
   if (!multi_content_) {
-    std::vector<AssociatedContentDriver*> drivers{tab_delegate};
+    std::vector<AssociatedContentDriver*> drivers{delegate};
     auto driver =
         std::make_unique<MultiAssociatedContentDriver>(std::move(drivers));
     SetMultiAssociatedContentDelegate(std::move(driver));
   } else {
-    multi_content_->AddContent(tab_delegate);
+    multi_content_->AddContent(delegate);
   }
 
   OnAssociatedContentInfoChanged();
   MaybeSeedOrClearSuggestions();
   MaybeFetchOrClearContentStagedConversation();
 }
 
-void ConversationHandler::RemoveAssociatedTab(mojom::AvailableTabPtr tab) {
+void ConversationHandler::RemoveAssociation(AssociatedContentDriver* delegate) {
   if (HasAnyHistory()) {
     return;
   }
@@ -581,8 +581,7 @@ void ConversationHandler::RemoveAssociatedTab(mojom::AvailableTabPtr tab) {
   }
 
   if (multi_content_) {
-    auto* tab_delegate = ai_chat_service_->GetAssociatedContent(tab);
-    multi_content_->RemoveContent(tab_delegate);
+    multi_content_->RemoveContent(delegate);
   }
 
   OnAssociatedContentInfoChanged();

components/ai_chat/core/browser/conversation_handler.h

@@ -58,6 +58,7 @@ class AIChatService;
 class AssociatedArchiveContent;
 class AIChatCredentialManager;
 class MultiAssociatedContentDriver;
+class AssociatedContentDriver;
 
 // Performs all conversation-related operations, responsible for sending
 // messages to the conversation engine, handling the responses, and owning
@@ -238,10 +239,11 @@ class ConversationHandler : public mojom::ConversationHandler,
   const mojom::Model& GetCurrentModel();
   const std::vector<mojom::ConversationTurnPtr>& GetConversationHistory() const;
 
+  void AddAssociation(AssociatedContentDriver* delegate);
+  void RemoveAssociation(AssociatedContentDriver* delegate);
+
   // mojom::ConversationHandler
   void GetState(GetStateCallback callback) override;
-  void AddAssociatedTab(mojom::AvailableTabPtr tab) override;
-  void RemoveAssociatedTab(mojom::AvailableTabPtr tab) override;
   void GetConversationHistory(GetConversationHistoryCallback callback) override;
   void RateMessage(bool is_liked,
                    const std::string& turn_uuid,

components/ai_chat/core/common/mojom/ai_chat.mojom

@@ -383,6 +383,10 @@ interface AIChatUIHandler {
   ManagePremium();
   HandleVoiceRecognition(string conversation_uuid);
 
+  // Methods for managing multi-tab content association
+  AssociateTab(AvailableTab tab, string conversation_uuid);
+  DisassociateTab(AvailableTab tab, string conversation_uuid);
+
   // This might be a no-op if the UI isn't closeable
   CloseUI();
 
@@ -449,9 +453,6 @@ struct ConversationEntriesState {
 // Browser-side handler for a Conversation
 interface ConversationHandler {
   GetState() => (ConversationState conversation_state);
-  AddAssociatedTab(AvailableTab tab);
-  RemoveAssociatedTab(AvailableTab tab);
-
   GetConversationUuid() => (string conversation_uuid);
 
   // Get all the possible models for any conversation

components/ai_chat/resources/page/components/site_titles/index.tsx

@@ -58,7 +58,7 @@ function SiteTitle(props: SiteTitleProps) {
           {props.detail.title}
         </p>
       </div>
-      {removable && <Button fab kind='plain-faint' onClick={() => conversation.conversationHandler?.removeAssociatedTab(tab)}>
+      {removable && <Button fab kind='plain-faint' onClick={() => aiChat.uiHandler?.disassociateTab(tab, conversation.conversationUuid!)}>
         <Icon name='trash' />
       </Button>}
     </div>