From 433782824e542c71f5e830d0c7e6c83604fe3f4e Mon Sep 17 00:00:00 2001 From: Pete Miller Date: Thu, 23 Oct 2025 11:00:09 -0700 Subject: [PATCH] [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 --- .../ai_chat/content_agent_tool_provider.cc | 4 ++ .../core/browser/conversation_handler.cc | 20 +++++++++ .../core/browser/conversation_handler.h | 13 ++++++ .../browser/conversation_handler_unittest.cc | 43 +++++++++++++++++++ ...ck_untrusted_conversation_handler_client.h | 1 + .../core/browser/tools/tool_provider.cc | 8 ++++ .../core/browser/tools/tool_provider.h | 16 +++++++ .../core/common/mojom/untrusted_frame.mojom | 4 ++ .../untrusted_conversation_frame_api.ts | 11 +++++ 9 files changed, 120 insertions(+) diff --git a/browser/ai_chat/content_agent_tool_provider.cc b/browser/ai_chat/content_agent_tool_provider.cc index bcfc17a8aea..e613bded280 100644 --- a/browser/ai_chat/content_agent_tool_provider.cc +++ b/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_, diff --git a/components/ai_chat/core/browser/conversation_handler.cc b/components/ai_chat/core/browser/conversation_handler.cc index 57624cdcded..055edfffac9 100644 --- a/components/ai_chat/core/browser/conversation_handler.cc +++ b/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& models = model_service_->GetModels(); auto default_model_key = model_service_->GetDefaultModelKey(); diff --git a/components/ai_chat/core/browser/conversation_handler.h b/components/ai_chat/core/browser/conversation_handler.h index ef87c9d2bee..a8173c95507 100644 --- a/components/ai_chat/core/browser/conversation_handler.h +++ b/components/ai_chat/core/browser/conversation_handler.h @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -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& get_task_tab_ids() const { return task_tab_ids_; } + void SetEngineForTesting(std::unique_ptr 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 task_tab_ids_; + raw_ptr ai_chat_service_; raw_ptr model_service_; raw_ptr credential_manager_; diff --git a/components/ai_chat/core/browser/conversation_handler_unittest.cc b/components/ai_chat/core/browser/conversation_handler_unittest.cc index 7a9456cc66b..8181ebbbcac 100644 --- a/components/ai_chat/core/browser/conversation_handler_unittest.cc +++ b/components/ai_chat/core/browser/conversation_handler_unittest.cc @@ -6,6 +6,7 @@ #include "brave/components/ai_chat/core/browser/conversation_handler.h" #include +#include #include #include #include @@ -97,6 +98,12 @@ class MockToolProvider : public ToolProvider { MOCK_METHOD(void, OnNewGenerationLoop, (), (override)); MOCK_METHOD(std::vector>, 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( + conversation_handler_->GetEngineForTesting()); + + // This test verifies that the conversation client is informed of the start + // of a content task from a ToolProvider. + NiceMock 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"); diff --git a/components/ai_chat/core/browser/mock_untrusted_conversation_handler_client.h b/components/ai_chat/core/browser/mock_untrusted_conversation_handler_client.h index 976294678cf..fef7a8b0654 100644 --- a/components/ai_chat/core/browser/mock_untrusted_conversation_handler_client.h +++ b/components/ai_chat/core/browser/mock_untrusted_conversation_handler_client.h @@ -41,6 +41,7 @@ class MockUntrustedConversationHandlerClient AssociatedContentChanged, (std::vector), (override)); + MOCK_METHOD(void, ContentTaskStarted, (int32_t), (override)); private: mojo::Receiver conversation_ui_receiver_{ diff --git a/components/ai_chat/core/browser/tools/tool_provider.cc b/components/ai_chat/core/browser/tools/tool_provider.cc index 0304517b239..2dad862fbc7 100644 --- a/components/ai_chat/core/browser/tools/tool_provider.cc +++ b/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 diff --git a/components/ai_chat/core/browser/tools/tool_provider.h b/components/ai_chat/core/browser/tools/tool_provider.h index e8434561313..ab8bc68427a 100644 --- a/components/ai_chat/core/browser/tools/tool_provider.h +++ b/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 #include #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 observers_; }; } // namespace ai_chat diff --git a/components/ai_chat/core/common/mojom/untrusted_frame.mojom b/components/ai_chat/core/common/mojom/untrusted_frame.mojom index 463aa7a77a2..6892da729e4 100644 --- a/components/ai_chat/core/common/mojom/untrusted_frame.mojom +++ b/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 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 diff --git a/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts b/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts index 5a9213f40dc..abda83f8c02 100644 --- a/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts +++ b/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 { + this.setPartialState({ contentTaskTabId: tabId }) + }, + ) + this.conversationObserver.onEntriesUIStateChanged.addListener( (state: Mojom.ConversationEntriesState) => { this.setPartialState(state)