Skip to content

Commit

Permalink
[AI Chat]: Add TabInformer so the front end can be aware of open tabs
Browse files Browse the repository at this point in the history
  • Loading branch information
fallaciousreasoning committed Jan 14, 2025
1 parent 54ef40b commit 9b8e446
Show file tree
Hide file tree
Showing 21 changed files with 390 additions and 19 deletions.
4 changes: 3 additions & 1 deletion browser/brave_content_browser_client.cc
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/settings_helper.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/tab_informer.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/untrusted_frame.mojom.h"
#include "brave/components/ai_rewriter/common/buildflags/buildflags.h"
#include "brave/components/body_sniffer/body_sniffer_throttle.h"
Expand Down Expand Up @@ -624,7 +625,8 @@ void BraveContentBrowserClient::RegisterWebUIInterfaceBrokers(
if (ai_chat::features::IsAIChatEnabled()) {
registry.ForWebUI<AIChatUI>()
.Add<ai_chat::mojom::AIChatUIHandler>()
.Add<ai_chat::mojom::Service>();
.Add<ai_chat::mojom::Service>()
.Add<ai_chat::mojom::TabInformer>();
registry.ForWebUI<AIChatUntrustedConversationUI>()
.Add<ai_chat::mojom::UntrustedUIHandler>()
.Add<ai_chat::mojom::UntrustedConversationHandler>();
Expand Down
1 change: 1 addition & 0 deletions browser/resources/settings/sources.gni
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,7 @@ brave_settings_ts_extra_deps =
brave_settings_mojo_files = [
"$root_gen_dir/brave/components/ai_chat/core/common/mojom/settings_helper.mojom-webui.ts",
"$root_gen_dir/brave/components/ai_chat/core/common/mojom/ai_chat.mojom-webui.ts",
"$root_gen_dir/brave/components/ai_chat/core/common/mojom/tab_informer.mojom-webui.ts",
"$root_gen_dir/brave/components/ai_chat/core/common/mojom/untrusted_frame.mojom-webui.ts",
]

Expand Down
2 changes: 2 additions & 0 deletions browser/ui/BUILD.gn
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ source_set("ui") {
# It doesn't make sense to view the webcompat webui on iOS & Android.
if (!is_android && !is_ios) {
sources += [
"ai_chat/tab_informer.cc",
"ai_chat/tab_informer.h",
"webui/webcompat_reporter/webcompat_reporter_ui.cc",
"webui/webcompat_reporter/webcompat_reporter_ui.h",
]
Expand Down
125 changes: 125 additions & 0 deletions browser/ui/ai_chat/tab_informer.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
// Copyright (c) 2025 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

#include "brave/browser/ui/ai_chat/tab_informer.h"

#include <algorithm>
#include <iterator>
#include <utility>

#include "base/strings/utf_string_conversions.h"
#include "brave/components/ai_chat/core/common/constants.h"
#include "brave/components/ai_chat/core/common/mojom/tab_informer.mojom.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/browser/ui/browser.h"
#include "chrome/browser/ui/browser_finder.h"
#include "chrome/browser/ui/browser_list.h"
#include "chrome/browser/ui/tabs/tab_model.h"
#include "content/public/browser/navigation_entry.h"
#include "content/public/browser/web_contents.h"
#include "mojo/public/cpp/bindings/receiver.h"

namespace ai_chat {

TabInformer::TabInformer(mojo::PendingReceiver<mojom::TabInformer> receiver,
content::WebContents* owner_contents)
: owner_contents_(owner_contents), receiver_(this, std::move(receiver)) {
tracker_.Init();
}

TabInformer::~TabInformer() = default;

content::WebContents* TabInformer::GetFromTab(const mojom::TabPtr& mojom_tab) {
const tabs::TabHandle handle = tabs::TabHandle(mojom_tab->id);
tabs::TabInterface* const tab = handle.Get();
if (!tab) {
return nullptr;
}
return tab->GetContents();
}

void TabInformer::AddListener(
mojo::PendingRemote<mojom::TabListener> listener) {
auto id = listeners_.Add(std::move(listener));

NotifyListener(listeners_.Get(id), GetState());
}

void TabInformer::OnTabStripModelChanged(
TabStripModel* tab_strip_model,
const TabStripModelChange& change,
const TabStripSelectionChange& selection) {
NotifyListeners();
}

void TabInformer::TabChangedAt(content::WebContents* contents,
int index,
TabChangeType change_type) {
NotifyListeners();
}

bool TabInformer::ShouldTrackBrowser(Browser* browser) {
return browser->is_type_normal() &&
browser->profile() ==
Profile::FromBrowserContext(owner_contents_->GetBrowserContext());
}

TabInformer::State TabInformer::GetState() {
State state;

// Used to determine whether a given window is active.
auto* active = chrome::FindLastActive();

for (Browser* browser : *BrowserList::GetInstance()) {
if (!ShouldTrackBrowser(browser)) {
continue;
}
auto window = mojom::Window::New();
window->is_active = browser == active;

auto* tab_strip_model = browser->tab_strip_model();
for (int i = 0; i < tab_strip_model->count(); ++i) {
auto* contents = tab_strip_model->GetWebContentsAt(i);
if (!contents) {
continue;
}

if (!ai_chat::kAllowedContentSchemes.contains(
contents->GetLastCommittedURL().scheme())) {
continue;
}

auto tab_state = mojom::Tab::New();
tab_state->id = tab_strip_model->GetTabHandleAt(i).raw_value();
tab_state->title = base::UTF16ToUTF8(contents->GetTitle());
tab_state->url = contents->GetVisibleURL();
tab_state->content_id =
contents->GetController().GetVisibleEntry()->GetUniqueID();
window->tabs.push_back(std::move(tab_state));
}

state.push_back(std::move(window));
}
return state;
}

void TabInformer::NotifyListeners() {
if (listeners_.empty()) {
return;
}
auto state = GetState();
for (auto& listener : listeners_) {
State clone;
std::ranges::transform(state, std::back_inserter(clone),
[](auto& window) { return window.Clone(); });
NotifyListener(listener.get(), std::move(clone));
}
}

void TabInformer::NotifyListener(mojom::TabListener* listener, State state) {
listener->TabsChanged(std::move(state));
}

} // namespace ai_chat
66 changes: 66 additions & 0 deletions browser/ui/ai_chat/tab_informer.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
// Copyright (c) 2025 The Brave Authors. All rights reserved.
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

#ifndef BRAVE_BROWSER_UI_AI_CHAT_TAB_INFORMER_H_
#define BRAVE_BROWSER_UI_AI_CHAT_TAB_INFORMER_H_

#include <string_view>
#include <vector>

#include "base/memory/raw_ptr.h"
#include "brave/components/ai_chat/core/common/mojom/tab_informer.mojom-forward.h"
#include "brave/components/ai_chat/core/common/mojom/tab_informer.mojom.h"
#include "chrome/browser/ui/browser_tab_strip_tracker.h"
#include "chrome/browser/ui/browser_tab_strip_tracker_delegate.h"
#include "chrome/browser/ui/tabs/tab_strip_model_observer.h"
#include "content/public/browser/web_contents.h"
#include "mojo/public/cpp/bindings/pending_receiver.h"
#include "mojo/public/cpp/bindings/receiver.h"
#include "mojo/public/cpp/bindings/remote_set.h"

namespace ai_chat {

class TabInformer : public mojom::TabInformer,
public BrowserTabStripTrackerDelegate,
public TabStripModelObserver {
public:
using State = std::vector<mojom::WindowPtr>;
// Note: |TabInformer| should not outlive |owner_contents_|.
TabInformer(mojo::PendingReceiver<mojom::TabInformer> receiver,
content::WebContents* owner_contents);
~TabInformer() override;

static content::WebContents* GetFromTab(const mojom::TabPtr& tab);

// mojom::TabInformer
void AddListener(mojo::PendingRemote<mojom::TabListener> listener) override;

// TabStripModelObserver:
void OnTabStripModelChanged(
TabStripModel* tab_strip_model,
const TabStripModelChange& change,
const TabStripSelectionChange& selection) override;
void TabChangedAt(content::WebContents* contents,
int index,
TabChangeType change_type) override;

// BrowserTabStripTrackerDelegate:
bool ShouldTrackBrowser(Browser* browser) override;

private:
State GetState();
void NotifyListeners();
void NotifyListener(mojom::TabListener* listener, State state);

raw_ptr<content::WebContents> owner_contents_;

mojo::Receiver<mojom::TabInformer> receiver_;
mojo::RemoteSet<mojom::TabListener> listeners_;
BrowserTabStripTracker tracker_{this, this};
};

} // namespace ai_chat

#endif // BRAVE_BROWSER_UI_AI_CHAT_TAB_INFORMER_H_
10 changes: 10 additions & 0 deletions browser/ui/webui/ai_chat/ai_chat_ui.cc
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

#include "brave/browser/ui/webui/ai_chat/ai_chat_ui.h"

#include <memory>
#include <utility>

#include "brave/browser/ai_chat/ai_chat_service_factory.h"
Expand All @@ -14,6 +15,7 @@
#include "brave/components/ai_chat/core/browser/ai_chat_service.h"
#include "brave/components/ai_chat/core/browser/constants.h"
#include "brave/components/ai_chat/core/browser/utils.h"
#include "brave/components/ai_chat/core/common/constants.h"
#include "brave/components/ai_chat/core/common/features.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "brave/components/ai_chat/core/common/pref_names.h"
Expand Down Expand Up @@ -162,6 +164,14 @@ void AIChatUI::BindInterface(
std::move(parent_ui_frame_receiver));
}

void AIChatUI::BindInterface(
mojo::PendingReceiver<ai_chat::mojom::TabInformer> tab_informer_receiver) {
#if !BUILDFLAG(IS_ANDROID)
tab_informer_ = std::make_unique<ai_chat::TabInformer>(
std::move(tab_informer_receiver), web_ui()->GetWebContents());
#endif
}

bool AIChatUIConfig::IsWebUIEnabled(content::BrowserContext* browser_context) {
return ai_chat::IsAIChatEnabled(
user_prefs::UserPrefs::Get(browser_context)) &&
Expand Down
7 changes: 7 additions & 0 deletions browser/ui/webui/ai_chat/ai_chat_ui.h
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,14 @@

#include "brave/browser/ui/webui/ai_chat/ai_chat_ui_page_handler.h"
#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/tab_informer.mojom.h"
#include "chrome/browser/ui/webui/top_chrome/top_chrome_web_ui_controller.h"
#include "content/public/browser/web_ui_controller.h"
#include "ui/webui/mojo_web_ui_controller.h"
#include "ui/webui/untrusted_web_ui_controller.h"

#if !BUILDFLAG(IS_ANDROID)
#include "brave/browser/ui/ai_chat/tab_informer.h"
#include "chrome/browser/ui/webui/top_chrome/top_chrome_webui_config.h"
#else
#include "content/public/browser/webui_config.h"
Expand All @@ -40,6 +42,8 @@ class AIChatUI : public ui::MojoWebUIController {
void BindInterface(mojo::PendingReceiver<ai_chat::mojom::Service> receiver);
void BindInterface(mojo::PendingReceiver<ai_chat::mojom::ParentUIFrame>
parent_ui_frame_receiver);
void BindInterface(
mojo::PendingReceiver<ai_chat::mojom::TabInformer> tab_informer_receiver);

// Set by WebUIContentsWrapperT. TopChromeWebUIController provides default
// implementation for this but we don't use it.
Expand All @@ -52,6 +56,9 @@ class AIChatUI : public ui::MojoWebUIController {

private:
std::unique_ptr<ai_chat::AIChatUIPageHandler> page_handler_;
#if !BUILDFLAG(IS_ANDROID)
std::unique_ptr<ai_chat::TabInformer> tab_informer_;
#endif

base::WeakPtr<TopChromeWebUIController::Embedder> embedder_;
raw_ptr<Profile> profile_ = nullptr;
Expand Down
68 changes: 68 additions & 0 deletions browser/ui/webui/ai_chat/ai_chat_ui_page_handler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,11 @@
#include <utility>
#include <vector>

#include "base/functional/bind.h"
#include "base/functional/callback_forward.h"
#include "brave/browser/ai_chat/ai_chat_service_factory.h"
#include "brave/browser/ai_chat/ai_chat_urls.h"
#include "brave/browser/ui/ai_chat/tab_informer.h"
#include "brave/browser/ui/side_panel/ai_chat/ai_chat_side_panel_utils.h"
#include "brave/components/ai_chat/core/browser/ai_chat_service.h"
#include "brave/components/ai_chat/core/browser/constants.h"
Expand All @@ -25,6 +28,7 @@
#include "chrome/browser/ui/singleton_tabs.h"
#include "components/favicon/core/favicon_service.h"
#include "content/public/browser/browser_context.h"
#include "content/public/browser/navigation_entry.h"
#include "content/public/browser/storage_partition.h"
#include "content/public/browser/web_contents.h"
#include "content/public/browser/web_contents_observer.h"
Expand Down Expand Up @@ -52,6 +56,46 @@ constexpr char kURLManagePremium[] = "https://account.brave.com/";

namespace ai_chat {

namespace {

class WaitForCommit : public content::WebContentsObserver {
public:
WaitForCommit(
content::WebContents* contents,
base::OnceCallback<void(content::WebContents* contents)> on_loaded)
: WebContentsObserver(contents), on_loaded_(std::move(on_loaded)) {}
~WaitForCommit() override = default;

void DidFinishNavigation(content::NavigationHandle* handle) override {
if (handle->IsInMainFrame() && handle->HasCommitted()) {
std::move(on_loaded_).Run(web_contents());
delete this;
}
}

void WebContentsDestroyed() override { delete this; }

private:
base::OnceCallback<void(content::WebContents* contents)> on_loaded_;
};

// Note: After session restore we need to ensure the WebContents is loaded
// before associating content with a conversation.
void EnsureWebContentsLoaded(
content::WebContents* contents,
base::OnceCallback<void(content::WebContents* contents)> on_loaded) {
if (!contents->GetController().NeedsReload()) {
std::move(on_loaded).Run(contents);
return;
}

// Deletes when the load completes or the WebContents is destroyed
new WaitForCommit(contents, std::move(on_loaded));
contents->GetController().LoadIfNecessary();
}

} // namespace

using mojom::CharacterType;
using mojom::ConversationTurn;
using mojom::ConversationTurnVisibility;
Expand Down Expand Up @@ -234,6 +278,30 @@ void AIChatUIPageHandler::BindRelatedConversation(
conversation->Bind(std::move(receiver), std::move(conversation_ui_handler));
}

void AIChatUIPageHandler::AssociateTab(mojom::TabPtr tab,
const std::string& conversation_uuid) {
auto* contents = ai_chat::TabInformer::GetFromTab(tab);
if (!contents) {
return;
}

EnsureWebContentsLoaded(
contents, base::BindOnce(
[](const std::string& conversation_uuid,
content::WebContents* contents) {
auto* tab_helper =
ai_chat::AIChatTabHelper::FromWebContents(contents);
if (!tab_helper) {
return;
}

AIChatServiceFactory::GetForBrowserContext(
contents->GetBrowserContext())
->AssociateContent(tab_helper, conversation_uuid);
},
conversation_uuid));
}

void AIChatUIPageHandler::NewConversation(
mojo::PendingReceiver<mojom::ConversationHandler> receiver,
mojo::PendingRemote<mojom::ConversationUI> conversation_ui_handler) {
Expand Down
2 changes: 2 additions & 0 deletions browser/ui/webui/ai_chat/ai_chat_ui_page_handler.h
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ class AIChatUIPageHandler : public mojom::AIChatUIHandler,
mojo::PendingReceiver<mojom::ConversationHandler> receiver,
mojo::PendingRemote<mojom::ConversationUI> conversation_ui_handler)
override;
void AssociateTab(mojom::TabPtr tab,
const std::string& conversation_uuid) override;
void NewConversation(
mojo::PendingReceiver<mojom::ConversationHandler> receiver,
mojo::PendingRemote<mojom::ConversationUI> conversation_ui_handler)
Expand Down
Loading

0 comments on commit 9b8e446

Please sign in to comment.