diff --git a/browser/history_embeddings/README.md b/browser/history_embeddings/README.md index cf85dc52830..86d1c6af8e6 100644 --- a/browser/history_embeddings/README.md +++ b/browser/history_embeddings/README.md @@ -75,10 +75,10 @@ without touching the upstream header (see implementation of `passage_embeddings::mojom::PassageEmbeddingsService`. Exposes a direct `BindPassageEmbedder(receiver, model_files, cb)` entry point used by the controller; constructs a - `BraveBatchPassageEmbedder` around the supplied files. Also hosts the - static `WebContents*` → bind-callback registry used by - `UntrustedLocalAIUI::BindInterface` (the registered callback routes - to the active `BraveBatchPassageEmbedder`). + `BraveBatchPassageEmbedder` around the supplied files. Also exposes + `BindLocalAIReceiver(...)` which the controller forwards to from + `UntrustedLocalAIUI::BindInterface` so the WASM page can register its + `PassageEmbedderFactory`. - **`brave_batch_passage_embedder.{h,cc}`** — In-process implementation of `passage_embeddings::mojom::PassageEmbedder` and diff --git a/browser/history_embeddings/brave_passage_embeddings_service.cc b/browser/history_embeddings/brave_passage_embeddings_service.cc index 4f40be2a2eb..8b4cd8523cf 100644 --- a/browser/history_embeddings/brave_passage_embeddings_service.cc +++ b/browser/history_embeddings/brave_passage_embeddings_service.cc @@ -8,29 +8,13 @@ #include #include "base/check.h" -#include "base/containers/flat_map.h" #include "base/feature_list.h" #include "base/functional/bind.h" #include "base/logging.h" -#include "base/no_destructor.h" #include "components/history_embeddings/core/history_embeddings_features.h" -#include "content/public/browser/web_contents.h" namespace passage_embeddings { -namespace { - -using BindRegistry = - base::flat_map; - -BindRegistry& GetBindRegistry() { - static base::NoDestructor registry; - return *registry; -} - -} // namespace - BravePassageEmbeddingsService::BravePassageEmbeddingsService( BackgroundWebContentsFactory background_web_contents_factory) : background_web_contents_factory_( @@ -40,27 +24,10 @@ BravePassageEmbeddingsService::BravePassageEmbeddingsService( BravePassageEmbeddingsService::~BravePassageEmbeddingsService() = default; -// static -void BravePassageEmbeddingsService::SetBindCallbackForWebContents( - content::WebContents* web_contents, - BindCallback callback) { - GetBindRegistry()[web_contents] = std::move(callback); -} - -// static -void BravePassageEmbeddingsService::RemoveBindCallbackForWebContents( - content::WebContents* web_contents) { - GetBindRegistry().erase(web_contents); -} - -// static -void BravePassageEmbeddingsService::BindForWebContents( - content::WebContents* web_contents, +void BravePassageEmbeddingsService::BindLocalAIReceiver( mojo::PendingReceiver receiver) { - auto& registry = GetBindRegistry(); - auto it = registry.find(web_contents); - if (it != registry.end()) { - it->second.Run(std::move(receiver)); + if (batch_embedder_) { + batch_embedder_->BindLocalAIReceiver(std::move(receiver)); } } diff --git a/browser/history_embeddings/brave_passage_embeddings_service.h b/browser/history_embeddings/brave_passage_embeddings_service.h index 60ec5ef8559..d07a31630f2 100644 --- a/browser/history_embeddings/brave_passage_embeddings_service.h +++ b/browser/history_embeddings/brave_passage_embeddings_service.h @@ -15,10 +15,6 @@ #include "mojo/public/cpp/bindings/pending_receiver.h" #include "services/passage_embeddings/public/mojom/passage_embeddings.mojom.h" -namespace content { -class WebContents; -} - namespace passage_embeddings { // In-process implementation of @@ -35,10 +31,6 @@ namespace passage_embeddings { // renderer-side lifecycle (background WebContents, LocalAIService // receiver set, PassageEmbedderFactory remote, and the // mojom::PassageEmbedder pipe back to the controller). -// -// Also exposes the static WebContents -> BindCallback registry used by -// UntrustedLocalAIUI::BindInterface to route renderer-side -// LocalAIService bindings to the active embedder. class BravePassageEmbeddingsService : public mojom::PassageEmbeddingsService { public: using BackgroundWebContentsFactory = @@ -52,19 +44,10 @@ class BravePassageEmbeddingsService : public mojom::PassageEmbeddingsService { BravePassageEmbeddingsService& operator=( const BravePassageEmbeddingsService&) = delete; - // Registry for routing mojo binding requests from the background - // WebContents (on guest OTR) back to the active embedder. - // UntrustedLocalAIUI::BindInterface calls BindForWebContents; the - // controller installs SetBindCallbackForWebContents when the - // embedder creates its background contents. - using BindCallback = base::RepeatingCallback)>; - static void SetBindCallbackForWebContents(content::WebContents* web_contents, - BindCallback callback); - static void RemoveBindCallbackForWebContents( - content::WebContents* web_contents); - static void BindForWebContents( - content::WebContents* web_contents, + // Forwards a renderer-side LocalAIService binding request to the + // active BatchEmbedder, if one exists. UntrustedLocalAIUI::BindInterface + // reaches this through the controller singleton. + void BindLocalAIReceiver( mojo::PendingReceiver receiver); // Direct in-process equivalent of mojom::PassageEmbeddingsService::LoadModels diff --git a/browser/history_embeddings/brave_passage_embeddings_service_controller.cc b/browser/history_embeddings/brave_passage_embeddings_service_controller.cc index 5d08bf4a56b..676c00cc6f0 100644 --- a/browser/history_embeddings/brave_passage_embeddings_service_controller.cc +++ b/browser/history_embeddings/brave_passage_embeddings_service_controller.cc @@ -46,12 +46,7 @@ mojom::PassagePriority ToMojom(PassagePriority priority) { } } -void InstallBindCallback(base::WeakPtr weak_embedder, - content::WebContents* web_contents) { - auto bind_cb = base::BindRepeating( - &BraveBatchPassageEmbedder::BindLocalAIReceiver, weak_embedder); - BravePassageEmbeddingsService::SetBindCallbackForWebContents( - web_contents, std::move(bind_cb)); +void TagWebContentsForTaskManager(content::WebContents* web_contents) { task_manager::WebContentsTags::CreateForToolContents( web_contents, IDS_LOCAL_AI_TASK_MANAGER_TITLE); } @@ -71,7 +66,7 @@ void OnGuestProfileCreated( BravePassageEmbeddingsServiceController::Get()->ObserveGuestOTRProfile(otr); auto contents = std::make_unique( otr, GURL(local_ai::kUntrustedLocalAIURL), weak_embedder.get(), - base::BindOnce(&InstallBindCallback, weak_embedder)); + base::BindOnce(&TagWebContentsForTaskManager)); std::move(callback).Run(std::move(contents)); } @@ -193,6 +188,13 @@ void BravePassageEmbeddingsServiceController::ResetServiceRemote() { otr_profile_observation_.Reset(); } +void BravePassageEmbeddingsServiceController::BindLocalAIReceiver( + mojo::PendingReceiver receiver) { + if (service_) { + service_->BindLocalAIReceiver(std::move(receiver)); + } +} + void BravePassageEmbeddingsServiceController::ObserveGuestOTRProfile( Profile* otr_profile) { CHECK(otr_profile); diff --git a/browser/history_embeddings/brave_passage_embeddings_service_controller.h b/browser/history_embeddings/brave_passage_embeddings_service_controller.h index 9cbd160edcc..a41a7c005b7 100644 --- a/browser/history_embeddings/brave_passage_embeddings_service_controller.h +++ b/browser/history_embeddings/brave_passage_embeddings_service_controller.h @@ -74,6 +74,12 @@ class BravePassageEmbeddingsServiceController // BrowserContextImpl's `rph_with_bc_reference` NOTREACHED). void ObserveGuestOTRProfile(Profile* otr_profile); + // Routes a renderer-side LocalAIService binding from + // UntrustedLocalAIUI::BindInterface to the active embedder. No-op + // when no service is alive. + void BindLocalAIReceiver( + mojo::PendingReceiver receiver); + private: friend class base::NoDestructor; diff --git a/browser/history_embeddings/brave_passage_embeddings_service_unittest.cc b/browser/history_embeddings/brave_passage_embeddings_service_unittest.cc index 9a19ce4da51..942ae7bd33a 100644 --- a/browser/history_embeddings/brave_passage_embeddings_service_unittest.cc +++ b/browser/history_embeddings/brave_passage_embeddings_service_unittest.cc @@ -11,7 +11,6 @@ #include "base/functional/bind.h" #include "base/memory/raw_ptr.h" -#include "base/test/bind.h" #include "base/test/run_until.h" #include "base/test/scoped_feature_list.h" #include "base/test/task_environment.h" @@ -299,32 +298,32 @@ TEST_F(BravePassageEmbeddingsServiceTest, BindAgainAfterEmbedderRemoteReset) { EXPECT_TRUE(load2->load_success.Get()); } -TEST_F(BravePassageEmbeddingsServiceTest, BindRegistryRoutesToService) { - // Manually exercise the static registry: install a bind callback for - // a fake WebContents pointer, call BindForWebContents, and confirm - // the callback fires. - auto* fake_web_contents = - reinterpret_cast(uintptr_t{0xdeadbeef}); - bool invoked = false; - BravePassageEmbeddingsService::SetBindCallbackForWebContents( - fake_web_contents, - base::BindLambdaForTesting( - [&](mojo::PendingReceiver) { - invoked = true; - })); +TEST_F(BravePassageEmbeddingsServiceTest, BindLocalAIReceiverNoopWithoutBatch) { + // Service-level forwarder is a no-op when no BatchEmbedder is bound. + // The receiver is dropped; the test just confirms the call doesn't + // crash and the remote sees a disconnect. + mojo::Remote remote; + service_->BindLocalAIReceiver(remote.BindNewPipeAndPassReceiver()); + base::test::TestFuture disconnected; + remote.set_disconnect_handler(disconnected.GetCallback()); + EXPECT_TRUE(disconnected.Wait()); +} - mojo::PendingRemote dummy_remote; - BravePassageEmbeddingsService::BindForWebContents( - fake_web_contents, dummy_remote.InitWithNewPipeAndPassReceiver()); - EXPECT_TRUE(invoked); +TEST_F(BravePassageEmbeddingsServiceTest, BindLocalAIReceiverForwardsToBatch) { + // Once a BatchEmbedder is alive, the forwarder hands the receiver to + // it — RegisterPassageEmbedderFactory then drives the load to Ready. + auto load = IssueLoad(); + ASSERT_TRUE(last_created_web_contents_); + ASSERT_TRUE(last_delegate_); - BravePassageEmbeddingsService::RemoveBindCallbackForWebContents( - fake_web_contents); - invoked = false; - mojo::PendingRemote dummy_remote2; - BravePassageEmbeddingsService::BindForWebContents( - fake_web_contents, dummy_remote2.InitWithNewPipeAndPassReceiver()); - EXPECT_FALSE(invoked); + mojo::Remote local_ai_remote; + service_->BindLocalAIReceiver(local_ai_remote.BindNewPipeAndPassReceiver()); + local_ai_remote->RegisterPassageEmbedderFactory(fake_factory_.BindRemote()); + local_ai_remote.FlushForTesting(); + + ASSERT_TRUE( + base::test::RunUntil([&] { return fake_factory_.init_count() > 0; })); + EXPECT_TRUE(load->load_success.Get()); } } // namespace passage_embeddings diff --git a/browser/ui/webui/local_ai/local_ai_ui.cc b/browser/ui/webui/local_ai/local_ai_ui.cc index d7d40c271aa..a61a420fab2 100644 --- a/browser/ui/webui/local_ai/local_ai_ui.cc +++ b/browser/ui/webui/local_ai/local_ai_ui.cc @@ -8,7 +8,7 @@ #include #include -#include "brave/browser/history_embeddings/brave_passage_embeddings_service.h" +#include "brave/browser/history_embeddings/brave_passage_embeddings_service_controller.h" #include "brave/components/local_ai/core/local_ai.mojom.h" #include "brave/components/local_ai/core/url_constants.h" #include "brave/components/local_ai/resources/grit/candle_embedding_module_generated.h" @@ -48,8 +48,8 @@ WEB_UI_CONTROLLER_TYPE_IMPL(UntrustedLocalAIUI) void UntrustedLocalAIUI::BindInterface( mojo::PendingReceiver receiver) { - passage_embeddings::BravePassageEmbeddingsService::BindForWebContents( - web_ui()->GetWebContents(), std::move(receiver)); + passage_embeddings::BravePassageEmbeddingsServiceController::Get() + ->BindLocalAIReceiver(std::move(receiver)); } ///////////////////////////////////////////////////////////////////////////////