From 80f6a9e8bb7a35ca00eb502805d872c1ef3609cf Mon Sep 17 00:00:00 2001
From: Mike Wittman <[email protected]>
Date: Thu, 21 May 2026 18:03:04 -0700
Subject: [PATCH] [embeddings] Fix BFCache race in embeddings attribution

PageEmbeddingsService maintained an internal reference to the active
page for a WebContents which was not cleared upon navigation. If a
navigation occurred and the previous page entered the Back/Forward
Cache, any pending embedding computation for that page could complete
and be incorrectly attributed to the new navigation's metadata in
HistoryEmbeddingsService.

This CL fixes the issue by overriding PrimaryPageChanged in the
WebContentsEventsObserver to clear the active page reference and reset
the embedding state. This ensures that computations for BFCached or
stale pages are correctly identified as no longer being the active page
for the WebContents.

Fixed: 514039947
Change-Id: Ic22432007c35c27aca27c5c5b4b4853934b2a0fd
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7856826
Reviewed-by: Zekun Jiang <[email protected]>
Commit-Queue: Mike Wittman <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1634639}
---

diff --git a/components/page_content_annotations/content/page_embeddings_service.cc b/components/page_content_annotations/content/page_embeddings_service.cc
index 2faae13..5bdd6180 100644
--- a/components/page_content_annotations/content/page_embeddings_service.cc
+++ b/components/page_content_annotations/content/page_embeddings_service.cc
@@ -89,6 +89,19 @@
     }
   }
 
+  void PrimaryPageChanged(content::Page& page) override {
+    auto loc =
+        page_embeddings_service_->web_contents_states_.find(web_contents());
+    if (loc != page_embeddings_service_->web_contents_states_.end()) {
+      if (auto* computing =
+              std::get_if<Computing>(&loc->second.embeddings_state)) {
+        page_embeddings_service_->embedder_->TryCancel(computing->task_id);
+      }
+      loc->second.page = nullptr;
+      loc->second.embeddings_state = Unavailable{};
+    }
+  }
+
   void WebContentsDestroyed() override {
     page_embeddings_service_->web_contents_states_.erase(web_contents());
   }
diff --git a/components/page_content_annotations/content/page_embeddings_service_unittest.cc b/components/page_content_annotations/content/page_embeddings_service_unittest.cc
index be3fb33..078ff63 100644
--- a/components/page_content_annotations/content/page_embeddings_service_unittest.cc
+++ b/components/page_content_annotations/content/page_embeddings_service_unittest.cc
@@ -1167,4 +1167,62 @@
       base::MakeRefCounted<RefCountedPDFText>("pdf text content"));
 }
 
+// Validates that embeddings computed for a page that is no longer the primary
+// page (e.g. it was navigated away from but is still alive in BFCache) are
+// ignored and do not notify observers.
+TEST_F(PageEmbeddingsServiceTest, BFCacheRaceReproduction) {
+  std::unique_ptr<content::WebContents> web_contents =
+      CreateTestWebContentsWithVisibility(content::Visibility::HIDDEN);
+
+  ObserverMock observer;
+  EXPECT_CALL(observer, GetDefaultPriority)
+      .WillRepeatedly(Return(PageEmbeddingsService::kDefault));
+  EXPECT_CALL(observer, GetUsageMode)
+      .WillRepeatedly(Return(PageEmbeddingsService::kOnDemand));
+  page_embeddings_service().AddObserver(&observer);
+
+  passage_embeddings::Embedder::ComputePassagesEmbeddingsCallback
+      compute_passages_embeddings_callback;
+
+  EXPECT_CALL(embedder_mock(), ComputePassagesEmbeddings)
+      .WillOnce(
+          [&](passage_embeddings::PassagePriority priority,
+              std::vector<std::string> passages,
+              passage_embeddings::Embedder::ComputePassagesEmbeddingsCallback
+                  callback) {
+            compute_passages_embeddings_callback = std::move(callback);
+            return 1;
+          });
+
+  // 1. Initial page load (attacker.com).
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents.get(), GURL("https://attacker.com"));
+  content::Page& page1 = web_contents->GetPrimaryPage();
+  base::WeakPtr<content::Page> page1_weak = page1.GetWeakPtr();
+
+  // 2. Content extracted for page 1.
+  page_embeddings_service().OnPageContentExtracted(
+      page1, base::MakeRefCounted<RefCountedAnnotatedPageContent>());
+
+  ASSERT_FALSE(compute_passages_embeddings_callback.is_null());
+
+  // 3. Navigate to page 2 (victim.com).
+  EXPECT_CALL(embedder_mock(), TryCancel(1));
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents.get(), GURL("https://victim.com"));
+  ASSERT_NE(nullptr, page1_weak)
+      << "Page 1 was destroyed upon navigation. BFCache simulation failed.";
+
+  // 4. Complete embedding for page 1.
+  // OnPageEmbeddingsAvailable should NOT be called because
+  // PrimaryPageChanged cleared the state.
+  EXPECT_CALL(observer, OnPageEmbeddingsAvailable(testing::_)).Times(0);
+
+  std::move(compute_passages_embeddings_callback)
+      .Run({"passage"}, {passage_embeddings::Embedding({1.0f})}, 1,
+           passage_embeddings::ComputeEmbeddingsStatus::kSuccess);
+
+  page_embeddings_service().RemoveObserver(&observer);
+}
+
 }  // namespace page_content_annotations

diff --git a/components/page_content_annotations/content/page_embeddings_service_unittest.cc b/components/page_content_annotations/content/page_embeddings_service_unittest.cc
index be3fb33..078ff63 100644
--- a/components/page_content_annotations/content/page_embeddings_service_unittest.cc
+++ b/components/page_content_annotations/content/page_embeddings_service_unittest.cc
@@ -1167,4 +1167,62 @@
       base::MakeRefCounted<RefCountedPDFText>("pdf text content"));
 }
 
+// Validates that embeddings computed for a page that is no longer the primary
+// page (e.g. it was navigated away from but is still alive in BFCache) are
+// ignored and do not notify observers.
+TEST_F(PageEmbeddingsServiceTest, BFCacheRaceReproduction) {
+  std::unique_ptr<content::WebContents> web_contents =
+      CreateTestWebContentsWithVisibility(content::Visibility::HIDDEN);
+
+  ObserverMock observer;
+  EXPECT_CALL(observer, GetDefaultPriority)
+      .WillRepeatedly(Return(PageEmbeddingsService::kDefault));
+  EXPECT_CALL(observer, GetUsageMode)
+      .WillRepeatedly(Return(PageEmbeddingsService::kOnDemand));
+  page_embeddings_service().AddObserver(&observer);
+
+  passage_embeddings::Embedder::ComputePassagesEmbeddingsCallback
+      compute_passages_embeddings_callback;
+
+  EXPECT_CALL(embedder_mock(), ComputePassagesEmbeddings)
+      .WillOnce(
+          [&](passage_embeddings::PassagePriority priority,
+              std::vector<std::string> passages,
+              passage_embeddings::Embedder::ComputePassagesEmbeddingsCallback
+                  callback) {
+            compute_passages_embeddings_callback = std::move(callback);
+            return 1;
+          });
+
+  // 1. Initial page load (attacker.com).
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents.get(), GURL("https://attacker.com"));
+  content::Page& page1 = web_contents->GetPrimaryPage();
+  base::WeakPtr<content::Page> page1_weak = page1.GetWeakPtr();
+
+  // 2. Content extracted for page 1.
+  page_embeddings_service().OnPageContentExtracted(
+      page1, base::MakeRefCounted<RefCountedAnnotatedPageContent>());
+
+  ASSERT_FALSE(compute_passages_embeddings_callback.is_null());
+
+  // 3. Navigate to page 2 (victim.com).
+  EXPECT_CALL(embedder_mock(), TryCancel(1));
+  content::NavigationSimulator::NavigateAndCommitFromBrowser(
+      web_contents.get(), GURL("https://victim.com"));
+  ASSERT_NE(nullptr, page1_weak)
+      << "Page 1 was destroyed upon navigation. BFCache simulation failed.";
+
+  // 4. Complete embedding for page 1.
+  // OnPageEmbeddingsAvailable should NOT be called because
+  // PrimaryPageChanged cleared the state.
+  EXPECT_CALL(observer, OnPageEmbeddingsAvailable(testing::_)).Times(0);
+
+  std::move(compute_passages_embeddings_callback)
+      .Run({"passage"}, {passage_embeddings::Embedding({1.0f})}, 1,
+           passage_embeddings::ComputeEmbeddingsStatus::kSuccess);
+
+  page_embeddings_service().RemoveObserver(&observer);
+}
+
 }  // namespace page_content_annotations