Drop WebContents → BindCallback registry on BravePassageEmbeddingsSevice (#36458)

Drop WebContents → BindCallback registry on BravePassageEmbeddingsService

There was ever at most one active embedder/BG WebContents pair, so the
static map was 0- or 1-entry overhead. Replace it with a direct
forwarder:

* BravePassageEmbeddingsService::BindLocalAIReceiver hands a
  PendingReceiver<LocalAIService> to batch_embedder_ if alive.
* BravePassageEmbeddingsServiceController exposes the same as an
  instance method, accessible via Get().
* UntrustedLocalAIUI::BindInterface routes through the controller.

Eliminates SetBindCallbackForWebContents (only producer call site went
away) and RemoveBindCallbackForWebContents (was unused in production —
the registry leaked one entry per BG WebContents lifecycle).

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

browser/history_embeddings/brave_passage_embeddings_service.cc

@@ -8,29 +8,13 @@
 #include <utility>
 
 #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<content::WebContents*,
-                   BravePassageEmbeddingsService::BindCallback>;
-
-BindRegistry& GetBindRegistry() {
-  static base::NoDestructor<BindRegistry> 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<local_ai::mojom::LocalAIService> 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));
   }
 }
 

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<void(
-      mojo::PendingReceiver<local_ai::mojom::LocalAIService>)>;
-  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<local_ai::mojom::LocalAIService> receiver);
 
   // Direct in-process equivalent of mojom::PassageEmbeddingsService::LoadModels

browser/history_embeddings/brave_passage_embeddings_service_controller.cc

@@ -46,12 +46,7 @@ mojom::PassagePriority ToMojom(PassagePriority priority) {
   }
 }
 
-void InstallBindCallback(base::WeakPtr<BraveBatchPassageEmbedder> 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<local_ai::BackgroundWebContentsImpl>(
       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<local_ai::mojom::LocalAIService> receiver) {
+  if (service_) {
+    service_->BindLocalAIReceiver(std::move(receiver));
+  }
+}
+
 void BravePassageEmbeddingsServiceController::ObserveGuestOTRProfile(
     Profile* otr_profile) {
   CHECK(otr_profile);

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<local_ai::mojom::LocalAIService> receiver);
+
  private:
   friend class base::NoDestructor<BravePassageEmbeddingsServiceController>;
 

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<content::WebContents*>(uintptr_t{0xdeadbeef});
-  bool invoked = false;
-  BravePassageEmbeddingsService::SetBindCallbackForWebContents(
-      fake_web_contents,
-      base::BindLambdaForTesting(
-          [&](mojo::PendingReceiver<local_ai::mojom::LocalAIService>) {
-            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<local_ai::mojom::LocalAIService> remote;
+  service_->BindLocalAIReceiver(remote.BindNewPipeAndPassReceiver());
+  base::test::TestFuture<void> disconnected;
+  remote.set_disconnect_handler(disconnected.GetCallback());
+  EXPECT_TRUE(disconnected.Wait());
+}
 
-  mojo::PendingRemote<local_ai::mojom::LocalAIService> 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<local_ai::mojom::LocalAIService> dummy_remote2;
-  BravePassageEmbeddingsService::BindForWebContents(
-      fake_web_contents, dummy_remote2.InitWithNewPipeAndPassReceiver());
-  EXPECT_FALSE(invoked);
+  mojo::Remote<local_ai::mojom::LocalAIService> 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

browser/ui/webui/local_ai/local_ai_ui.cc

@@ -8,7 +8,7 @@
 #include <memory>
 #include <utility>
 
-#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<mojom::LocalAIService> receiver) {
-  passage_embeddings::BravePassageEmbeddingsService::BindForWebContents(
-      web_ui()->GetWebContents(), std::move(receiver));
+  passage_embeddings::BravePassageEmbeddingsServiceController::Get()
+      ->BindLocalAIReceiver(std::move(receiver));
 }
 
 ///////////////////////////////////////////////////////////////////////////////