From 68be1396f12e2ddf65a274f819c92842b92bc3c7 Mon Sep 17 00:00:00 2001 From: Pete Miller Date: Wed, 3 Jun 2026 00:06:01 -0700 Subject: [PATCH] [AI Chat] Tab Management Tool - close action (#36894) * [AI Chat] Tab Management Tool - close action Adds the "close" action to the tab management tool, reimplemented on top of the refactored handlers. Closing is asynchronous (beforeunload handlers can delay WebContents destruction), so a self-owned TabsClosedWaiter waits for the tabs to be destroyed, or a fallback timeout, before reporting the resulting tab list. Duplicate tab_ids are de-duplicated and indices are re-queried per close to avoid acting on a stale or destroyed WebContents. --- browser/ai_chat/tools/tab_management_tool.cc | 186 ++++++++++++++++++ browser/ai_chat/tools/tab_management_tool.h | 1 + .../tools/tab_management_tool_browsertest.cc | 47 +++++ .../tools/tab_management_tool_unittest.cc | 5 + 4 files changed, 239 insertions(+) diff --git a/browser/ai_chat/tools/tab_management_tool.cc b/browser/ai_chat/tools/tab_management_tool.cc index c83ab3d35d6..91619a7a130 100644 --- a/browser/ai_chat/tools/tab_management_tool.cc +++ b/browser/ai_chat/tools/tab_management_tool.cc @@ -16,15 +16,22 @@ #include "base/containers/adapters.h" #include "base/containers/fixed_flat_map.h" #include "base/containers/flat_map.h" +#include "base/containers/flat_set.h" #include "base/functional/bind.h" +#include "base/functional/callback.h" #include "base/json/json_reader.h" #include "base/json/json_writer.h" #include "base/logging.h" #include "base/memory/raw_ptr.h" +#include "base/memory/raw_ref.h" +#include "base/memory/weak_ptr.h" #include "base/notreached.h" +#include "base/sequence_checker.h" #include "base/strings/string_number_conversions.h" #include "base/strings/utf_string_conversions.h" #include "base/task/sequenced_task_runner.h" +#include "base/time/time.h" +#include "base/timer/timer.h" #include "base/values.h" #include "brave/components/ai_chat/core/browser/tools/tool_input_properties.h" #include "brave/components/ai_chat/core/browser/tools/tool_utils.h" @@ -35,6 +42,7 @@ #include "chrome/browser/ui/browser.h" #include "chrome/browser/ui/browser_window/public/browser_window_interface.h" #include "chrome/browser/ui/browser_window/public/browser_window_interface_iterator.h" +#include "chrome/browser/ui/tabs/tab_enums.h" #include "chrome/browser/ui/tabs/tab_group_model.h" #include "chrome/browser/ui/tabs/tab_model.h" #include "chrome/browser/ui/tabs/tab_strip_model.h" @@ -44,6 +52,7 @@ #include "components/tabs/public/tab_group.h" #include "components/tabs/public/tab_interface.h" #include "content/public/browser/web_contents.h" +#include "content/public/browser/web_contents_observer.h" #include "ui/base/base_window.h" static_assert(BUILDFLAG(ENABLE_AI_CHAT_TAB_MANAGEMENT_TOOL)); @@ -53,6 +62,12 @@ namespace ai_chat { namespace { +// Fallback timeout for tab removal operations. Tab closure can involve user +// interaction via unload (beforeunload) handlers, so the WebContents may +// outlive the close request. If all tabs are removed within the timeout we +// report the post-close state; otherwise we report whatever tabs remain. +constexpr base::TimeDelta kTabRemovalTimeout = base::Seconds(10); + // Returns a sorted, de-duplicated list of indices that are valid for the // provided TabStripModel. Any indices outside of the current tab bounds are // dropped. @@ -319,6 +334,110 @@ int MoveResolvedTabsToStrip(const std::vector& tabs_to_move, return static_cast(moved_indices.size()); } +// Waits for a set of tabs to be destroyed after a close request, then runs +// |on_done|. Tab closing is asynchronous: CloseWebContentsAt() initiates +// closure but WebContents destruction can be delayed (e.g. by a beforeunload +// handler), so reporting immediately would still list the tabs as present. +// The callback runs once all tabs are gone or the fallback timeout fires, +// whichever comes first. Self-owned: it deletes itself after completion. +class TabsClosedWaiter { + public: + static void Run(std::vector handles, + base::OnceClosure on_done, + base::TimeDelta fallback_timeout = kTabRemovalTimeout) { + (new TabsClosedWaiter(std::move(handles), std::move(on_done), + fallback_timeout)) + ->Start(); + } + + private: + // Observes a single WebContents and notifies the waiter when it is + // destroyed. The waiter never dereferences the WebContents, so observing + // through destruction is safe. + class TabObserver : public content::WebContentsObserver { + public: + TabObserver(TabsClosedWaiter& waiter, content::WebContents* web_contents) + : content::WebContentsObserver(web_contents), waiter_(waiter) {} + + void WebContentsDestroyed() override { waiter_->CheckAndMaybeFinish(); } + + private: + const raw_ref waiter_; + }; + + TabsClosedWaiter(std::vector handles, + base::OnceClosure on_done, + base::TimeDelta timeout) + : handles_(std::move(handles)), + on_done_(std::move(on_done)), + timeout_(timeout) {} + + void Start() { + DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); + if (handles_.empty()) { + ForceFinish(); + return; + } + tab_observers_.reserve(handles_.size()); + for (tabs::TabHandle handle : handles_) { + if (auto* tab = handle.Get()) { + tab_observers_.push_back( + std::make_unique(*this, tab->GetContents())); + } + } + + // Some tabs may already be gone (e.g. closed synchronously); check on the + // next task so the caller's callback is never invoked re-entrantly. + base::SequencedTaskRunner::GetCurrentDefault()->PostTask( + FROM_HERE, base::BindOnce(&TabsClosedWaiter::CheckAndMaybeFinish, + weak_factory_.GetWeakPtr())); + timer_.Start(FROM_HERE, timeout_, + base::BindOnce(&TabsClosedWaiter::ForceFinish, + weak_factory_.GetWeakPtr())); + } + + void CheckAndMaybeFinish() { + DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); + // Keep waiting while any of the tabs is still alive. + if (std::ranges::any_of(handles_, [](tabs::TabHandle handle) { + return handle.Get() != nullptr; + })) { + return; + } + ForceFinish(); + } + + void ForceFinish() { + DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_); + if (finished_) { + return; + } + finished_ = true; + timer_.Stop(); + + // Post the callback on the next task, matching the rest of the tool, so any + // resulting active-window change has settled before the tab list is read. + if (on_done_) { + base::SequencedTaskRunner::GetCurrentDefault()->PostTask( + FROM_HERE, std::move(on_done_)); + } + + // Self-delete on the next task. Deferring (rather than clearing + // |tab_observers_| now) avoids destroying the TabObserver whose + // WebContentsDestroyed() callback is currently on the stack. + base::SequencedTaskRunner::GetCurrentDefault()->DeleteSoon(FROM_HERE, this); + } + + std::vector handles_; + std::vector> tab_observers_; + base::OnceClosure on_done_; + base::OneShotTimer timer_; + const base::TimeDelta timeout_; + bool finished_ = false; + SEQUENCE_CHECKER(sequence_checker_); + base::WeakPtrFactory weak_factory_{this}; +}; + } // namespace TabManagementTool::TabManagementTool(Profile* profile) : profile_(profile) {} @@ -503,6 +622,8 @@ void TabManagementTool::UseTool(const std::string& input_json, HandleMoveTabs(std::move(callback), dict); } else if (*action == "move_group") { HandleMoveGroup(std::move(callback), dict); + } else if (*action == "close") { + HandleCloseTabs(std::move(callback), dict); } else if (*action == "create_group") { HandleCreateGroup(std::move(callback), dict); } else if (*action == "update_group") { @@ -954,6 +1075,71 @@ void TabManagementTool::HandleMoveGroup(UseToolCallback callback, } } +void TabManagementTool::HandleCloseTabs(UseToolCallback callback, + const base::DictValue& params) { + const auto* tab_ids = params.FindList("tab_ids"); + + if (!tab_ids || tab_ids->empty()) { + std::move(callback).Run( + CreateContentBlocksForText( + "Missing or empty 'tab_ids' array for close operation"), + {}); + return; + } + + // Resolve the requested tabs to unique handles, in request order. We must not + // retain the ResolvedTab entries across the close loop below: they hold + // raw_ptrs to the TabInterface/WebContents, and CloseWebContentsAt() can + // destroy a WebContents synchronously, which would leave those raw_ptrs + // dangling. The ResolvedTab vector is therefore consumed (and destroyed) + // here, before any tab is closed. + std::vector handles; + base::flat_set seen_handles; + for (const ResolvedTab& tab : + ResolveTabs(HandleIdsFromList(*tab_ids), profile_, + /*require_grouped=*/false)) { + tabs::TabHandle handle = tab.tab->GetHandle(); + if (seen_handles.insert(handle.raw_value()).second) { + handles.push_back(handle); + } + } + + std::vector handles_to_wait_for; + handles_to_wait_for.reserve(handles.size()); + for (tabs::TabHandle handle : handles) { + tabs::TabInterface* tab = handle.Get(); + if (!tab) { + continue; + } + BrowserWindowInterface* browser = tab->GetBrowserWindowInterface(); + if (!browser) { + continue; + } + TabStripModel* tab_strip = browser->GetTabStripModel(); + // Re-query the index against the current strip each time: closing an + // earlier tab in the same strip shifts the indices of the remaining tabs. + int index = tab_strip->GetIndexOfWebContents(tab->GetContents()); + if (index == TabStripModel::kNoTab) { + continue; + } + tab_strip->CloseWebContentsAt(index, TabCloseTypes::CLOSE_USER_GESTURE); + handles_to_wait_for.push_back(handle); + } + const size_t closed_count = handles_to_wait_for.size(); + + base::DictValue result; + result.Set("message", "Successfully closed " + + base::NumberToString(closed_count) + " tab(s)"); + + // Tab closure is asynchronous, so wait for the tabs to actually be destroyed + // (or time out) before reporting the resulting tab list. + TabsClosedWaiter::Run( + std::move(handles_to_wait_for), + base::BindOnce(&TabManagementTool::PostTaskSendResultWithTabList, + weak_ptr_factory_.GetWeakPtr(), std::move(callback), + std::move(result), std::nullopt)); +} + void TabManagementTool::HandleCreateGroup(UseToolCallback callback, const base::DictValue& params) { const auto* tab_ids = params.FindList("tab_ids"); diff --git a/browser/ai_chat/tools/tab_management_tool.h b/browser/ai_chat/tools/tab_management_tool.h index b74c667ab62..c46a6559b6e 100644 --- a/browser/ai_chat/tools/tab_management_tool.h +++ b/browser/ai_chat/tools/tab_management_tool.h @@ -46,6 +46,7 @@ class TabManagementTool : public Tool { void HandleListTabs(UseToolCallback callback); void HandleMoveTabs(UseToolCallback callback, const base::DictValue& params); void HandleMoveGroup(UseToolCallback callback, const base::DictValue& params); + void HandleCloseTabs(UseToolCallback callback, const base::DictValue& params); void HandleCreateGroup(UseToolCallback callback, const base::DictValue& params); void HandleUpdateGroup(UseToolCallback callback, diff --git a/browser/ai_chat/tools/tab_management_tool_browsertest.cc b/browser/ai_chat/tools/tab_management_tool_browsertest.cc index 0329062238e..5eebd314afb 100644 --- a/browser/ai_chat/tools/tab_management_tool_browsertest.cc +++ b/browser/ai_chat/tools/tab_management_tool_browsertest.cc @@ -109,6 +109,17 @@ size_t GetBrowserCount() { return count; } +size_t GetTabCount(Profile* profile) { + size_t count = 0; + for (BrowserWindowInterface* bwi : GetAllBrowserWindowInterfaces()) { + if (bwi->GetProfile() != profile) { + continue; + } + count += bwi->GetTabStripModel()->count(); + } + return count; +} + // Build a minimal expected windows skeleton for the current profile state so // tests can compare. We should always also test actual browser state to verify // the response has waited for the state to catch up with the commands. @@ -827,4 +838,40 @@ IN_PROC_BROWSER_TEST_F(TabManagementToolBrowserTest, } } +IN_PROC_BROWSER_TEST_F(TabManagementToolBrowserTest, CloseTabsAcrossWindows) { + TabManagementTool tool(profile()); + tool.UserPermissionGranted(""); + + Browser* b1 = browser(); + Browser* b2 = CreateBrowser(profile()); + + int t1 = AddTabAndGetHandle(b1, GURL("https://close1.test/")); + int t2 = AddTabAndGetHandle(b2, GURL("https://close2.test/")); + // A tab we keep open, to confirm only the requested tabs are closed. + int keep = AddTabAndGetHandle(b1, GURL("https://keep.test/")); + + ASSERT_NE(GetSessionIdForTabId(t1), GetSessionIdForTabId(t2)); + const size_t initial_tab_count = GetTabCount(profile()); + + std::string response = RunToolAndGetText(FROM_HERE, &tool, + absl::StrFormat( + R"JSON({ + "action": "close", + "tab_ids": [%d, %d] + })JSON", + t1, t2)); + + base::DictValue response_dict = base::test::ParseJsonDict(response); + EXPECT_THAT(response_dict, + base::test::IsSupersetOfValue(base::test::ParseJson( + R"JSON({"message":"Successfully closed 2 tab(s)"})JSON"))); + + // The waiter must not report until the tabs are actually destroyed. + EXPECT_FALSE(tabs::TabHandle(t1).Get()); + EXPECT_FALSE(tabs::TabHandle(t2).Get()); + EXPECT_TRUE(tabs::TabHandle(keep).Get()); + EXPECT_EQ(GetTabCount(profile()), initial_tab_count - 2); + ExpectOutputMatchesWindowSkeleton(FROM_HERE, response, profile()); +} + } // namespace ai_chat diff --git a/browser/ai_chat/tools/tab_management_tool_unittest.cc b/browser/ai_chat/tools/tab_management_tool_unittest.cc index a1807cc0ce0..f395a2fe220 100644 --- a/browser/ai_chat/tools/tab_management_tool_unittest.cc +++ b/browser/ai_chat/tools/tab_management_tool_unittest.cc @@ -154,6 +154,11 @@ TEST_F(TabManagementToolUnitTest, JsonAndArgumentValidationErrors) { RunTool(&tool, R"({"action":"move_group"})"), testing::HasSubstr("Missing 'group_id' for move_group operation")); + // close without tab_ids + EXPECT_THAT(RunTool(&tool, R"({"action":"close"})"), + testing::HasSubstr( + "Missing or empty 'tab_ids' array for close operation")); + // create_group without tab_ids EXPECT_THAT( RunTool(&tool, R"({"action":"create_group"})"),