[AXTreeFixing] Add AXTreeFixing service for ScreenAI and create in Router

This CL adds a new internal service router for the AXTree fixing
feature. This router connects directly to the ScreenAI service and will
be used as a single pipe from the BrowserContextKeyedService tree fixing
Router to the ScreenAI service. This is made internal to the tree_fixing
target, and only the AXTreeFixingServicesRouter can depend on it. In the
future there will be many such internal routers. The top-level Router
provides a public API for any client to perform an AXTree fixing
operation, such as IdentifyMainNode, which is added in this CL.

AX-Relnotes: N/A
Bug: 395134535, 401034541, 399384018
Change-Id: Ic8aa07f2884b6ab3bc61a13e1466bdf6ee0765c5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6325211
Reviewed-by: Nektarios Paisios <nektar@chromium.org>
Reviewed-by: Mustafa Emre Acer <meacer@chromium.org>
Commit-Queue: Mark Schillaci <mschillaci@google.com>
Cr-Commit-Position: refs/heads/main@{#1429730}

chrome/browser/accessibility/tree_fixing/BUILD.gn

@@ -17,4 +17,18 @@ if (!is_android) {
       "//content/public/browser:browser",
     ]
   }
+
+  # All clients that want to use AXTreeFixing need to do so via
+  # the router, and only the router can access the internal
+  # connections to downstream services.
+  source_set("internal") {
+    visibility = [ ":impl" ]
+
+    sources = [
+      "internal/ax_tree_fixing_screen_ai_service.cc",
+      "internal/ax_tree_fixing_screen_ai_service.h",
+    ]
+
+    deps = [ "//base" ]
+  }
 }

chrome/browser/accessibility/tree_fixing/ax_tree_fixing_services_router.cc

@@ -4,9 +4,53 @@
 
 #include "chrome/browser/accessibility/tree_fixing/ax_tree_fixing_services_router.h"
 
+#include <memory>
+
+#include "base/check.h"
+#include "chrome/browser/accessibility/tree_fixing/internal/ax_tree_fixing_screen_ai_service.h"
+#include "ui/accessibility/ax_tree_update.h"
+
 namespace tree_fixing {
 
-AXTreeFixingServicesRouter::AXTreeFixingServicesRouter(Profile* profile) {}
+AXTreeFixingServicesRouter::AXTreeFixingServicesRouter(Profile* profile)
+    : profile_(profile) {}
 AXTreeFixingServicesRouter::~AXTreeFixingServicesRouter() = default;
 
+void AXTreeFixingServicesRouter::IdentifyMainNode(
+    const ui::AXTreeUpdate& ax_tree,
+    MainNodeIdentificationCallback callback) {
+  // If this is the first time any client has requested tree fixing in a form
+  // that is handled by the ScreenAI service, then create an instance to connect
+  // to the service now.
+  if (!screen_ai_service_) {
+    screen_ai_service_ =
+        std::make_unique<AXTreeFixingScreenAIService>(this, profile_);
+  }
+
+  // Store the callback for later use, and make a request to ScreenAI.
+  pending_callbacks_.emplace_back(next_request_id_, std::move(callback));
+  screen_ai_service_->IdentifyMainNode(ax_tree, next_request_id_);
+  next_request_id_++;
+}
+
+void AXTreeFixingServicesRouter::OnMainNodeIdentified(ui::AXTreeID tree_id,
+                                                      ui::AXNodeID node_id,
+                                                      int request_id) {
+  CHECK(!pending_callbacks_.empty());
+
+  // Find the callback associated with the returned request ID, and call it with
+  // the identified tree_id and node_id for the upstream client to use. Remove
+  // the pending callback since we have fulfilled the contract.
+  for (auto it = pending_callbacks_.begin(); it != pending_callbacks_.end();
+       ++it) {
+    if (it->first == request_id) {
+      MainNodeIdentificationCallback callback = std::move(it->second);
+      pending_callbacks_.erase(it);
+      std::move(callback).Run(std::make_pair(tree_id, node_id));
+      return;
+    }
+  }
+  NOTREACHED();
+}
+
 }  // namespace tree_fixing

chrome/browser/accessibility/tree_fixing/ax_tree_fixing_services_router.h

@@ -5,16 +5,29 @@
 #ifndef CHROME_BROWSER_ACCESSIBILITY_TREE_FIXING_AX_TREE_FIXING_SERVICES_ROUTER_H_
 #define CHROME_BROWSER_ACCESSIBILITY_TREE_FIXING_AX_TREE_FIXING_SERVICES_ROUTER_H_
 
+#include <list>
+#include <memory>
+
+#include "chrome/browser/accessibility/tree_fixing/internal/ax_tree_fixing_screen_ai_service.h"
 #include "components/keyed_service/core/keyed_service.h"
 
 class Profile;
 
+namespace ui {
+struct AXTreeUpdate;
+}  // namespace ui
+
 namespace tree_fixing {
 
+using MainNodeIdentificationCallback =
+    base::OnceCallback<void(std::pair<ui::AXTreeID, ui::AXNodeID>)>;
+
 // This class handles the communication between the browser process and any
 // downstream services used to fix the AXTree, such as: the optimization guide,
 // Screen2x, Aratea, etc.
-class AXTreeFixingServicesRouter : public KeyedService {
+class AXTreeFixingServicesRouter
+    : public KeyedService,
+      public AXTreeFixingScreenAIService::MainNodeIdentificationDelegate {
  public:
   explicit AXTreeFixingServicesRouter(Profile* profile);
   AXTreeFixingServicesRouter(const AXTreeFixingServicesRouter&) = delete;
@@ -22,7 +35,29 @@ class AXTreeFixingServicesRouter : public KeyedService {
       delete;
   ~AXTreeFixingServicesRouter() override;
 
-  // TODO(mschillaci): Stubbed. Fill in with mojom bindings.
+  // --- Public APIs for any request to fix an AXTree ---
+
+  // Identifies the main node of an AXTree, and asynchronously returns the
+  // identified node_id and its associated tree_id as a std::pair via the
+  // provided callback. The AXTreeUpdate that clients provide to this method
+  // should represent a full AXTree for the page in order to accurately identify
+  // a main node. The AXTree should not have an existing node with Role kMain.
+  void IdentifyMainNode(const ui::AXTreeUpdate& ax_tree,
+                        MainNodeIdentificationCallback callback);
+
+ private:
+  // AXTreeFixingScreenAIService::MainNodeIdentificationDelegate overrides:
+  void OnMainNodeIdentified(ui::AXTreeID tree_id,
+                            ui::AXNodeID node_id,
+                            int request_id) override;
+
+  // ScreenAI related objects: service instance, and a list of callbacks/ids.
+  std::unique_ptr<AXTreeFixingScreenAIService> screen_ai_service_;
+  std::list<std::pair<int, MainNodeIdentificationCallback>> pending_callbacks_;
+  int next_request_id_ = 0;
+
+  // The Profile for the KeyedService for this instance.
+  Profile* profile_;
 };
 
 }  // namespace tree_fixing

chrome/browser/accessibility/tree_fixing/internal/ax_tree_fixing_screen_ai_service.cc

@@ -0,0 +1,61 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "chrome/browser/accessibility/tree_fixing/internal/ax_tree_fixing_screen_ai_service.h"
+
+#include <utility>
+
+#include "base/notreached.h"
+#include "chrome/browser/screen_ai/screen_ai_service_router.h"
+#include "chrome/browser/screen_ai/screen_ai_service_router_factory.h"
+#include "ui/accessibility/ax_enums.mojom.h"
+#include "ui/accessibility/ax_node_data.h"
+
+namespace tree_fixing {
+
+AXTreeFixingScreenAIService::AXTreeFixingScreenAIService(
+    MainNodeIdentificationDelegate& delegate,
+    Profile* profile)
+    : main_node_identification_delegate_(delegate), profile_(profile) {}
+
+AXTreeFixingScreenAIService::~AXTreeFixingScreenAIService() = default;
+
+void AXTreeFixingScreenAIService::IdentifyMainNode(
+    const ui::AXTreeUpdate& ax_tree_update,
+    int request_id) {
+  // Clients should not be sending requests for trees that have a kMain node.
+  for (const ui::AXNodeData& node : ax_tree_update.nodes) {
+    if (node.role == ax::mojom::Role::kMain) {
+      NOTREACHED();
+    }
+  }
+
+  // If the remote to ScreenAI has not yet been bound, do so now.
+  // TODO(crbug.com/401308988): Handle cases like not ready or disconnects.
+  if (!screen_ai_service_.is_bound() || !screen_ai_service_.is_connected()) {
+    mojo::PendingReceiver<screen_ai::mojom::Screen2xMainContentExtractor>
+        receiver = screen_ai_service_.BindNewPipeAndPassReceiver();
+    screen_ai::ScreenAIServiceRouterFactory::GetForBrowserContext(profile_)
+        ->BindMainContentExtractor(std::move(receiver));
+    screen_ai_service_.reset_on_disconnect();
+  }
+
+  // Identify the main node using ScreenAI.
+  screen_ai_service_->IdentifyMainNode(
+      ax_tree_update,
+      base::BindOnce(&AXTreeFixingScreenAIService::
+                         ProcessScreenAIMainNodeIdentificationResult,
+                     weak_ptr_factory_.GetWeakPtr(), request_id));
+}
+
+void AXTreeFixingScreenAIService::ProcessScreenAIMainNodeIdentificationResult(
+    const ui::AXTreeID& tree_id,
+    int node_id,
+    int request_id) {
+  // TODO(crbug.com/399383663): Add metrics, internal logic, etc.
+  main_node_identification_delegate_->OnMainNodeIdentified(tree_id, node_id,
+                                                           request_id);
+}
+
+}  // namespace tree_fixing

chrome/browser/accessibility/tree_fixing/internal/ax_tree_fixing_screen_ai_service.h

@@ -0,0 +1,81 @@
+// Copyright 2025 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CHROME_BROWSER_ACCESSIBILITY_TREE_FIXING_INTERNAL_AX_TREE_FIXING_SCREEN_AI_SERVICE_H_
+#define CHROME_BROWSER_ACCESSIBILITY_TREE_FIXING_INTERNAL_AX_TREE_FIXING_SCREEN_AI_SERVICE_H_
+
+#include "base/memory/raw_ref.h"
+#include "base/memory/weak_ptr.h"
+#include "chrome/browser/profiles/profile.h"
+#include "mojo/public/cpp/bindings/remote.h"
+#include "services/screen_ai/public/mojom/screen_ai_service.mojom.h"
+
+namespace tree_fixing {
+
+class AXTreeFixingScreenAIService final {
+ public:
+  // Delegate for clients that want to perform main node identification.
+  class MainNodeIdentificationDelegate {
+   protected:
+    MainNodeIdentificationDelegate() = default;
+
+   public:
+    MainNodeIdentificationDelegate(const MainNodeIdentificationDelegate&) =
+        delete;
+    MainNodeIdentificationDelegate& operator=(
+        const MainNodeIdentificationDelegate&) = delete;
+    virtual ~MainNodeIdentificationDelegate() = default;
+
+    // This method is used to communicate to the delegate (owner) of this
+    // instance, the main node that was identified via the IdentifyMainNode
+    // method. When calling IdentifyMainNode, the client must provide a
+    // request_id, and this ID is passed back to the client along with the
+    // tree_id and node_id. The request_id allows clients to make multiple
+    // requests in parallel and uniquely identify each response. It is the
+    // responsibility of the client to handle the logic behind a request_id,
+    // this service simply passes the id through.
+    virtual void OnMainNodeIdentified(ui::AXTreeID tree_id,
+                                      ui::AXNodeID node_id,
+                                      int request_id) = 0;
+  };
+
+  AXTreeFixingScreenAIService(MainNodeIdentificationDelegate& delegate,
+                              Profile* profile);
+  AXTreeFixingScreenAIService(const AXTreeFixingScreenAIService&) = delete;
+  AXTreeFixingScreenAIService& operator=(const AXTreeFixingScreenAIService&) =
+      delete;
+  ~AXTreeFixingScreenAIService();
+
+  // --- Public APIs for upstream clients (e.g. AXTreeFixingServicesRouter) ---
+
+  // Identifies the main node of an AXTreeUpdate. The client should provide a
+  // request_id, which is returned to the client along with a tree_id and
+  // node_id via a call to OnMainNodeIdentified.
+  void IdentifyMainNode(const ui::AXTreeUpdate& ax_tree_update, int request_id);
+
+ private:
+  // Internal method that processes results from the ScreenAI service before
+  // returning the results to the owner of this instance via the provided
+  // delegate.
+  void ProcessScreenAIMainNodeIdentificationResult(const ui::AXTreeID& tree_id,
+                                                   int node_id,
+                                                   int request_id);
+
+  // Delegate provided by client to receive main node identification results.
+  // Use a raw_ref since we do not own the delegate or control its lifecycle.
+  raw_ref<MainNodeIdentificationDelegate> main_node_identification_delegate_;
+
+  // Profile for the KeyedService that owns us.
+  raw_ptr<Profile> profile_;
+
+  // The remote of the ScreenAI service, the receiver is in a utility process.
+  mojo::Remote<screen_ai::mojom::Screen2xMainContentExtractor>
+      screen_ai_service_;
+
+  base::WeakPtrFactory<AXTreeFixingScreenAIService> weak_ptr_factory_{this};
+};
+
+}  // namespace tree_fixing
+
+#endif  // CHROME_BROWSER_ACCESSIBILITY_TREE_FIXING_INTERNAL_AX_TREE_FIXING_SCREEN_AI_SERVICE_H_

content/renderer/accessibility/annotations/ax_main_node_annotator_browsertest.cc

@@ -53,6 +53,9 @@ class MockMainNodeAnnotationService
     std::move(callback).Run(main_);
   }
 
+  void IdentifyMainNode(const ui::AXTreeUpdate& snapshot,
+                        IdentifyMainNodeCallback callback) override {}
+
   // Tests should not modify entries in these lists.
   std::vector<ui::AXNodeID> content_nodes_;
   ui::AXNodeID main_ = ui::kInvalidAXNodeID;

services/screen_ai/public/mojom/screen_ai_service.mojom

@@ -137,7 +137,15 @@ interface Screen2xMainContentExtractor {
 
   // Receives the accessibility tree as a snapshot, schedules processing, and
   // returns the main node id of the given tree.
+  // DEPRECATED, use IdentifyMainNode method below.
+  // TODO(crbug.com/401052200): Clean up existing calls and delete method.
   ExtractMainNode(ax.mojom.AXTreeUpdate snapshot) => (int32 main_node_id);
+
+  // Receives an AXTreeUpdate and identifies the main node. The AXTreeUpdate
+  // should represent an entire AXTree; not a diff or partial tree, in order
+  // to identify an accurate main node for the entire page.
+  IdentifyMainNode(ax.mojom.AXTreeUpdate ax_tree) =>
+    (ax.mojom.AXTreeID tree_id, int32 node_id);
 };
 
 // Provides an interface to the OCR functionality of the Screen AI service.

services/screen_ai/screen_ai_service_impl.cc

@@ -455,6 +455,20 @@ void ScreenAIService::ExtractMainNode(const ui::AXTreeUpdate& snapshot,
   }
 }
 
+void ScreenAIService::IdentifyMainNode(const ui::AXTreeUpdate& snapshot,
+                                       IdentifyMainNodeCallback callback) {
+  ui::AXTree tree;
+  std::optional<std::vector<int32_t>> content_node_ids;
+  bool success = ExtractMainContentInternal(snapshot, tree, content_node_ids);
+
+  if (success) {
+    ui::AXNodeID main_node_id = ComputeMainNode(&tree, *content_node_ids);
+    std::move(callback).Run(tree.GetAXTreeID(), main_node_id);
+  } else {
+    std::move(callback).Run(ui::AXTreeIDUnknown(), ui::kInvalidAXNodeID);
+  }
+}
+
 bool ScreenAIService::ExtractMainContentInternal(
     const ui::AXTreeUpdate& snapshot,
     ui::AXTree& tree,

services/screen_ai/screen_ai_service_impl.h

@@ -76,6 +76,8 @@ class ScreenAIService : public mojom::ScreenAIServiceFactory,
                           ExtractMainContentCallback callback) override;
   void ExtractMainNode(const ui::AXTreeUpdate& snapshot,
                        ExtractMainNodeCallback callback) override;
+  void IdentifyMainNode(const ui::AXTreeUpdate& snapshot,
+                        IdentifyMainNodeCallback callback) override;
 
   // mojom::ScreenAIServiceFactory:
   void InitializeMainContentExtraction(