[AI Chat] [Content Agent] Allow a ConversationHandler and UI to observe tab-related tasks from a ToolProvider (#31945)

* [AI Chat] [Content Agent] Allow a ConversationHandler and UI to observe tab-related tasks from a ToolProvider

This will allow the UI to know about the tab being operated on. That can be used to e.g. show a thumbnail of the tab, or ask the browser to make the tab active

browser/ai_chat/content_agent_tool_provider.cc

@@ -104,6 +104,10 @@ void ContentAgentToolProvider::GetOrCreateTabHandleForTask(
 
     task_tab_handle_ =
         tabs::TabInterface::GetFromContents(new_contents)->GetHandle();
+
+    for (auto& observer : observers_) {
+      observer.OnContentTaskStarted(task_tab_handle_.raw_value());
+    }
   }
   actor_service_->GetTask(task_id_)->AddTab(
       task_tab_handle_,

components/ai_chat/core/browser/conversation_handler.cc

@@ -141,6 +141,11 @@ ConversationHandler::ConversationHandler(
     conversation_capability_ = mojom::ConversationCapability::CONTENT_AGENT;
   }
 
+  // Observe tool providers
+  for (const auto& tool_provider : tool_providers_) {
+    tool_provider->AddObserver(this);
+  }
+
   // When a client disconnects, let observers know
   receivers_.set_disconnect_handler(
       base::BindRepeating(&ConversationHandler::OnClientConnectionChanged,
@@ -172,6 +177,9 @@ ConversationHandler::ConversationHandler(
 
 ConversationHandler::~ConversationHandler() {
   OnConversationDeleted();
+  for (const auto& tool_provider : tool_providers_) {
+    tool_provider->RemoveObserver(this);
+  }
 }
 
 void ConversationHandler::AddObserver(Observer* observer) {
@@ -1595,6 +1603,18 @@ void ConversationHandler::OnModelRemoved(const std::string& removed_key) {
   InitEngine();
 }
 
+void ConversationHandler::OnContentTaskStarted(int32_t tab_id) {
+  // Store the tab_id so consumers can validate a tab is used by a tool
+  // in this conversation, or that a tab can locate its controlling
+  // conversation.
+  task_tab_ids_.insert(tab_id);
+
+  // Notify clients so they may display the relationship in UI
+  for (auto& client : untrusted_conversation_ui_handlers_) {
+    client->ContentTaskStarted(tab_id);
+  }
+}
+
 void ConversationHandler::OnModelDataChanged() {
   const std::vector<mojom::ModelPtr>& models = model_service_->GetModels();
   auto default_model_key = model_service_->GetDefaultModelKey();

components/ai_chat/core/browser/conversation_handler.h

@@ -9,6 +9,7 @@
 #include <cstdint>
 #include <memory>
 #include <optional>
+#include <set>
 #include <string>
 #include <string_view>
 #include <utility>
@@ -67,6 +68,7 @@ class AssociatedContentManager;
 class ConversationHandler : public mojom::ConversationHandler,
                             public mojom::UntrustedConversationHandler,
                             public ModelService::Observer,
+                            public ToolProvider::Observer,
                             public ConversationHandlerForMetrics {
  public:
   using GeneratedTextCallback =
@@ -255,6 +257,8 @@ class ConversationHandler : public mojom::ConversationHandler,
 
   mojom::APIError current_error() const override;
 
+  const std::set<int32_t>& get_task_tab_ids() const { return task_tab_ids_; }
+
   void SetEngineForTesting(std::unique_ptr<EngineConsumer> engine_for_testing) {
     engine_ = std::move(engine_for_testing);
   }
@@ -296,6 +300,9 @@ class ConversationHandler : public mojom::ConversationHandler,
                              const std::string& new_key) override;
   void OnModelRemoved(const std::string& removed_key) override;
 
+  // ToolProvider::Observer
+  void OnContentTaskStarted(int32_t tab_id) override;
+
  private:
   friend class ::AIChatUIBrowserTest;
   FRIEND_TEST_ALL_PREFIXES(AIChatServiceUnitTest, DeleteAssociatedWebContent);
@@ -453,6 +460,12 @@ class ConversationHandler : public mojom::ConversationHandler,
   mojom::ConversationCapability conversation_capability_ =
       mojom::ConversationCapability::CHAT;
 
+  // Set of tab IDs that have been part of tasks whilst this conversation is
+  // in-memory. Since conversations are finite (limited by context size) and not
+  // held in memory forever, we don't currently prune the tab IDs once they
+  // close. Therefore, these are not guaranteed to be active.
+  std::set<int32_t> task_tab_ids_;
+
   raw_ptr<AIChatService, DanglingUntriaged> ai_chat_service_;
   raw_ptr<ModelService> model_service_;
   raw_ptr<AIChatCredentialManager, DanglingUntriaged> credential_manager_;

components/ai_chat/core/browser/conversation_handler_unittest.cc

@@ -6,6 +6,7 @@
 #include "brave/components/ai_chat/core/browser/conversation_handler.h"
 
 #include <algorithm>
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
@@ -97,6 +98,12 @@ class MockToolProvider : public ToolProvider {
   MOCK_METHOD(void, OnNewGenerationLoop, (), (override));
   MOCK_METHOD(std::vector<base::WeakPtr<Tool>>, GetTools, (), (override));
   MOCK_METHOD(void, StopAllTasks, (), (override));
+
+  void StartContentTask(int32_t tab_id) {
+    for (auto& observer : observers_) {
+      observer.OnContentTaskStarted(tab_id);
+    }
+  }
 };
 
 class MockConversationHandlerClient : public mojom::ConversationUI {
@@ -3461,6 +3468,42 @@ TEST_F(ConversationHandlerUnitTest, ToolUseEvents_ToolNotFound) {
   run_loop.Run();
 }
 
+TEST_F(ConversationHandlerUnitTest, ToolUseEvents_OnContentTaskStarted) {
+  conversation_handler_->associated_content_manager()->ClearContent();
+
+  int32_t test_tab_id = 1;
+
+  EXPECT_EQ(0u, conversation_handler_->get_task_tab_ids().size());
+
+  MockEngineConsumer* engine = static_cast<MockEngineConsumer*>(
+      conversation_handler_->GetEngineForTesting());
+
+  // This test verifies that the conversation client is informed of the start
+  // of a content task from a ToolProvider.
+  NiceMock<MockUntrustedConversationHandlerClient> untrusted_client(
+      conversation_handler_.get());
+  EXPECT_CALL(untrusted_client, ContentTaskStarted(test_tab_id));
+
+  base::RunLoop run_loop;
+  // Call to engine mocks the use tool request when the tool is first used.
+  // We do not need to complete the request as this test is verifying that
+  // the observation is made by the conversation client whilst the request
+  // is still in progress so that the UI may follow the progress of the action.
+  EXPECT_CALL(*engine, GenerateAssistantResponse)
+      .WillOnce(testing::WithArg<7>(
+          [&](EngineConsumer::GenerationDataCallback callback) {
+            mock_tool_provider_->StartContentTask(test_tab_id);
+            run_loop.QuitWhenIdle();  // QuitWhenIdle due to mojo connection
+          }));
+
+  // Submit a human entry to trigger the tool use
+  conversation_handler_->SubmitHumanConversationEntry(".", std::nullopt);
+  run_loop.Run();
+
+  EXPECT_EQ(1u, conversation_handler_->get_task_tab_ids().size());
+  EXPECT_EQ(test_tab_id, *conversation_handler_->get_task_tab_ids().begin());
+}
+
 TEST_F(ConversationHandlerUnitTest, AssociatingContentTriggersGetContent) {
   MockAssociatedContent content;
   content.SetTextContent("content");

components/ai_chat/core/browser/mock_untrusted_conversation_handler_client.h

@@ -41,6 +41,7 @@ class MockUntrustedConversationHandlerClient
               AssociatedContentChanged,
               (std::vector<mojom::AssociatedContentPtr>),
               (override));
+  MOCK_METHOD(void, ContentTaskStarted, (int32_t), (override));
 
  private:
   mojo::Receiver<mojom::UntrustedConversationUI> conversation_ui_receiver_{

components/ai_chat/core/browser/tools/tool_provider.cc

@@ -11,4 +11,12 @@ ToolProvider::ToolProvider() = default;
 
 ToolProvider::~ToolProvider() = default;
 
+void ToolProvider::AddObserver(Observer* observer) {
+  observers_.AddObserver(observer);
+}
+
+void ToolProvider::RemoveObserver(Observer* observer) {
+  observers_.RemoveObserver(observer);
+}
+
 }  // namespace ai_chat

components/ai_chat/core/browser/tools/tool_provider.h

@@ -6,9 +6,11 @@
 #ifndef BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_TOOL_PROVIDER_H_
 #define BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_TOOLS_TOOL_PROVIDER_H_
 
+#include <cstdint>
 #include <vector>
 
 #include "base/memory/weak_ptr.h"
+#include "base/observer_list.h"
 
 namespace ai_chat {
 
@@ -42,6 +44,17 @@ class ToolProvider {
   // but not a whole conversation.
   virtual void OnNewGenerationLoop() {}
 
+  class Observer : public base::CheckedObserver {
+   public:
+    ~Observer() override {}
+
+    // This ToolProvider has some Tool acting on a Tab
+    virtual void OnContentTaskStarted(int32_t tab_id) {}
+  };
+
+  void AddObserver(Observer* observer);
+  void RemoveObserver(Observer* observer);
+
   // Returns the list of tools available for the conversation.
   // The returned pointers *should* be valid as long as the ToolProvider exists
   // until either the ToolProvider is destroyed, or `OnNewGenerationLoop` is
@@ -56,6 +69,9 @@ class ToolProvider {
   // Attempts to stops all current tasks started by Tools from this
   // ToolProvider.
   virtual void StopAllTasks() {}
+
+ protected:
+  base::ObserverList<Observer> observers_;
 };
 
 }  // namespace ai_chat

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

@@ -77,6 +77,10 @@ interface UntrustedConversationUI {
   // Called when the associated content is changed - called when the untrusted
   // conversation frame is attached and when the content is changed.
   AssociatedContentChanged(array<AssociatedContent> associated_content);
+
+  // Called when an observable content-based task is started for any in-progress
+  // conversation entry.
+  ContentTaskStarted(int32 tab_id);
 };
 
 // UI-side handler for callbacks from UntrustedUIHandler

components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts

@@ -13,6 +13,11 @@ export type ConversationEntriesUIState = Mojom.ConversationEntriesState & {
   conversationHistory: Mojom.ConversationTurn[]
   isMobile: boolean
   associatedContent: Mojom.AssociatedContent[]
+  // TODO(https://github.com/brave/brave-browser/issues/49258):
+  // Store the tab ID of a task on the ToolUseEvent and not for the whole
+  // conversation, once multiple agentic tabs and tasks per conversation are
+  // supported.
+  contentTaskTabId?: number
 }
 
 // Default state before initial API call
@@ -134,6 +139,12 @@ export default class UntrustedConversationFrameAPI extends API<ConversationEntri
       },
     )
 
+    this.conversationObserver.contentTaskStarted.addListener(
+      (tabId: number) => {
+        this.setPartialState({ contentTaskTabId: tabId })
+      },
+    )
+
     this.conversationObserver.onEntriesUIStateChanged.addListener(
       (state: Mojom.ConversationEntriesState) => {
         this.setPartialState(state)