components/permissions/prediction_service/README.md

# Permissions Prediction Service

## Overview

The `//components/permissions/prediction_service` directory contains the core,
platform-agnostic engine for predicting user actions on permission prompts. Its
primary function is to provide a signal indicating the likelihood that a user
will grant a given permission request.

This component does **not** implement any user-facing UI itself. Instead, it
provides the underlying prediction logic that is consumed by the Chrome browser
layer (`//chrome/browser/permissions/prediction_service`) to make UI decisions,
such as choosing between a standard permission prompt and a quieter, less
intrusive UI.

## Usage

This component provides the core prediction logic but is not intended for direct
use by most features. The primary consumer is the
`PermissionsAiUiSelector` located in
`//chrome/browser/permissions/prediction_service/`.

The PredictionService is embedded in the broader flow of how a permission request
is handled in the browser and potentially presented to the user as a UI chip:

1. The `PermissionRequestManager` in `//components/permissions` manages the
   lifecycle of a permission prompt. It consults a list of registered
   `PermissionUiSelector` implementations to determine which UI to show.
1. `PermissionsAiUiSelector` is one such implementation. When its
   `SelectUiToUse` method is called, it determines which prediction model to
   use (server-side or one of several on-device models) based on feature flags
   and the availability of downloaded models.
1. It then gathers the necessary features by creating a
   `PredictionRequestFeatures` object. For some on-device models, this may also
   involve capturing a snapshot of the page or extracting its text content.
1. These features are passed to the appropriate handler in this component
   (e.g., `PredictionService` for remote predictions, or
   `PermissionsAiv4Handler` for on-device AI predictions). Here, the name
   prediction service is used specifically for the component that deals
   with the remote server-side model, but also as a name for the overall
   prediction-logic including the on-device models as well.
1. After the model runs, the handler returns a prediction (e.g., a grant
   likelihood or a relevance score).
1. The `PermissionsAiUiSelector` uses this prediction to make a
   final `Decision` on whether to show the standard permission prompt or a
   quieter UI, which is then passed back to the `PermissionRequestManager`.

Therefore, developers looking to understand how permission predictions
influence the UI should start by examining `PermissionsAiUiSelector`.

## Architecture

The main components of the prediction service engine are:

- **`PermissionUiSelector`**: The central abstraction that clients of the
  permission component, most often in the //chrome layer, implement. It defines
  the interface for deciding which UI to use for a permission request. The
  primary implementation using the prediction service logic is
  `PermissionsAiUiSelector` in
  `//chrome/browser/permissions/prediction_service/`.

- **`PredictionService`**: A `KeyedService` that handles making network
  requests to a remote prediction service. It takes a
  `PredictionRequestFeatures` object, serializes it into a protobuf, and sends
  it to the service endpoint.

- **On-Device Model Handlers**: For enhanced privacy and reduced latency, the
  service also utilizes local TFLite models. The downloading and execution of
  these models are handled by the `optimization_guide` component. More details
  are provided in the section about on-device models below.

- **`PredictionRequestFeatures`**: A struct that defines the set of features
  sent to the prediction service. These features include information about the
  user's past interactions with permission prompts (e.g., grant, deny, dismiss,
  and ignore counts) and the current request context (e.g., whether a user gesture
  was present). In workflows that use on-device AI models, it also includes the
  **`PermissionRequestRelevance`** score generated by those models. This score is
  the output of a multi-stage prediction process, as detailed in the
  "Model Stacking and Dependencies" section below. This entire feature set is then
  sent as input to the remote server-side model.

- **Protobufs**: The `prediction_service_messages.proto` file defines the
  request and response formats for communication with the remote prediction
  service.

### On-Device Models

The key classes within `//components/permissions/prediction_service` that implement the on-device logic are:

- **`PredictionModelHandler`**: Implements the `optimization_guide::ModelHandler` interface for the first-generation on-device model (CPSSv1). It observes the availability of the model file from the `OptimizationGuideModelProvider`. Upon request, it posts an inference task to a background thread using either `PredictionModelExecutor` or `PredictionSignatureModelExecutor`. The resulting grant likelihood is returned asynchronously to the caller via a `LookupResponseCallback`.

- **`PredictionModelExecutor`**: The executor for the first-generation model. Its `Preprocess` method converts the statistical features from a `GeneratePredictionsRequest` protobuf into the input tensors for the model. `Postprocess` converts the model's output tensor into a grant likelihood.

- **`PredictionSignatureModelExecutor`**: A variant of the above for the signature-based version of the statistical model. It maps feature names to the corresponding input tensors.

- **`PermissionsAiv3Handler`**: Implements the `optimization_guide::ModelHandler` interface for the AIv3 model. It observes model file availability from the `OptimizationGuideModelProvider`. When a prediction is requested, it posts an inference task to a background thread and enforces a timeout. The resulting `PermissionRequestRelevance` score is returned asynchronously to the caller via a `base::OnceCallback`.

- **`PermissionsAiv4Handler`**: Similar to its v3 counterpart, this handler orchestrates the execution of the AIv4 model and returns its `PermissionRequestRelevance` score via a `base::OnceCallback`. It includes additional logic to prevent concurrent model executions: if a new request arrives while one is in progress, the new request is immediately answered with an empty result to avoid UI delays. The in-progress request continues, but its result is discarded.

- **`PermissionsAiEncoderBase`**: A base class for the executors of the AIv3 and AIv4 models. It contains common logic for:

  - Resizing and converting a `SkBitmap` snapshot into a normalized input tensor.
  - Post-processing the model's raw float output into a `PermissionRequestRelevance` enum by comparing it against a set of thresholds.

- **`PermissionsAiv3Executor`**: The executor for the AIv3 model. Its `Preprocess` method prepares the snapshot input tensor. It can also parse model-specific thresholds from a `PermissionsAiv3ModelMetadata` protobuf.

- **`PermissionsAiv4Executor`**: The executor for the AIv4 model. Its `Preprocess` method prepares two input tensors: one for the snapshot and one for text embeddings (`passage_embeddings::Embedding`).

#### Model Comparison

The on-device models can be categorized into two main types: a statistical model (CPSSv1) and more advanced AI models (AIv3, AIv4). Each serves a different purpose and uses different inputs to generate predictions.

| Model Type | Key Inputs | Output | Model Flow |
| ---------- | -------------------------------------------------------------------------------------------------------- | ------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------ |
| **CPSSv1** | User interaction history (e.g., grant/deny rates), presence of a user gesture, and other client signals. | Grant Likelihood (`DiscretizedLikelihood`) | Standalone on-device model. Its output is a final decision. |
| **AIv3** | A snapshot of the web page content. | Permission Request Relevance | On-device model. Its output is used as an input feature for the server-side model. |
| **AIv4** | A snapshot of the web page content and text embeddings from the page's inner text. | Permission Request Relevance | On-device model. Depends on a text embedder model for one of its inputs. Its output is used as an input feature for the server-side model. |

The **`PermissionRequestRelevance`** output from the AIv3 and AIv4 models is an
enum that categorizes the relevance of a permission prompt into discrete
levels, from `kVeryLow` to `kVeryHigh`. This score is determined by comparing
the raw float output of the model against a set of predefined thresholds. A
higher relevance score suggests that the permission request is well-integrated
with the page's content and purpose, making it more likely to be a legitimate
and user-benefiting request.

#### Model Stacking and Dependencies

The prediction service employs a "model stacking" approach where the output of one model serves as an input to another. Specifically, the AIv3 and AIv4 on-device models act as the first stage in a two-stage process. They generate a `PermissionRequestRelevance` score, which is then combined with other statistical features and sent to the remote server-side model for a final grant likelihood prediction.

The AIv4 model itself has a preceding dependency: it requires text embeddings generated from the page content, which are provided by a separate passage embedder model. This creates a chain: `Passage Embedder -> AIv4 -> Server-side Model`. The passage embedder model is implemented in [`//services/passage_embeddings/`](https://source.chromium.org/chromium/chromium/src/+/main:services/passage_embeddings/).

#### Relevant Context

The on-device prediction models are delivered and executed through a set of classes that interact with the `//components/optimization_guide` component.

The model handlers and encoders of the prediction service build upon two main abstractions from the optimization guide:

- **`optimization_guide::ModelHandler`**: This is the primary interface for a feature that uses an on-device model. It is responsible for observing model updates from the `OptimizationGuideModelProvider`, managing the model's lifecycle, and posting inference tasks to a background thread. `PredictionModelHandler`, `PermissionsAiv3Handler`, and `PermissionsAiv4Handler` are all implementations of this class.

- **`optimization_guide::BaseModelExecutor`**: This class handles the low-level details of TFLite model execution. It loads the model file, runs the interpreter, and delegates the feature-specific logic of data transformation to its subclasses via `Preprocess` and `Postprocess` methods. Implementations in this component include `PredictionModelExecutor` and `PermissionsAiEncoderBase`.

  - **`optimization_guide::SignatureModelExecutor`**: A specialized version of `BaseModelExecutor` for models that use named input/output tensors (signatures) instead of indexed tensors. The `PredictionSignatureModelExecutor` is an implementation of this class.

## Testing

Unit tests for the remote `PredictionService` are located in `prediction_service_unittest.cc`. These tests use a `TestURLLoaderFactory` to mock network responses from the remote service.

Tests for the on-device model handlers can be found in:

- `permissions_aiv3_handler_unittest.cc`
- `permissions_aiv4_handler_unittest.cc`

These tests follow a common pattern:

- A `TestOptimizationGuideModelProvider` is used to simulate the delivery of a model file to the handler being tested.
- The tests use real, simple TFLite models (`//components/test/data/permissions/`) that are designed to return a constant, predictable value. This allows for testing the full execution flow.
- Fake model executors (e.g., `PermissionsAiv3EncoderFake`) are used to inject hooks that can inspect the input tensors and verify that data is being correctly processed and passed to the TFLite interpreter.
- Mock model handlers (e.g., `PermissionsAiv4HandlerMock`) are used to control the asynchronous execution flow, allowing for tests of timeout logic and the prevention of concurrent requests.
- `base::test::TestFuture` is used to wait for and capture the results from the asynchronous `ExecutionCallback`.

Finally, end-to-end integration tests can be found in `//chrome/browser/permissions/prediction_service/prediction_service_browsertest.cc`. These browser tests verify that the `PermissionsAiUiSelector` correctly uses the output of the prediction service (both remote and on-device) to make decisions about which permission UI to display.

## Relevant Context

To get a complete picture of the prediction service, consider the following key files for each major component:

### Permission UI Selector

- `//components/permissions/prediction_service/permission_ui_selector.h`
- `//components/permissions/prediction_service/permission_ui_selector.cc`

### Remote Prediction Service

- `//components/permissions/prediction_service/prediction_service.h`
- `//components/permissions/prediction_service/prediction_service.cc`
- `//components/permissions/prediction_service/prediction_service_base.h`
- `//components/permissions/prediction_service/prediction_request_features.h`
- `//components/permissions/prediction_service/prediction_common.h`
- `//components/permissions/prediction_service/prediction_common.cc`
- `//components/permissions/prediction_service/prediction_service_messages.proto`

### On-Device Statistical Model

- `//components/permissions/prediction_service/prediction_model_handler.h`
- `//components/permissions/prediction_service/prediction_model_handler.cc`
- `//components/permissions/prediction_service/prediction_model_executor.h`
- `//components/permissions/prediction_service/prediction_model_executor.cc`
- `//components/permissions/prediction_service/prediction_signature_model_executor.h`
- `//components/permissions/prediction_service/prediction_signature_model_executor.cc`
- `//components/permissions/prediction_service/prediction_model_metadata.proto`

### On-Device AI Model V3

- `//components/permissions/prediction_service/permissions_aiv3_handler.h`
- `//components/permissions/prediction_service/permissions_aiv3_handler.cc`
- `//components/permissions/prediction_service/permissions_aiv3_executor.h`
- `//components/permissions/prediction_service/permissions_aiv3_executor.cc`
- `//components/permissions/prediction_service/permissions_aiv3_model_metadata.proto`

### On-Device AI Model V4

- `//components/permissions/prediction_service/permissions_aiv4_handler.h`
- `//components/permissions/prediction_service/permissions_aiv4_handler.cc`
- `//components/permissions/prediction_service/permissions_aiv4_executor.h`
- `//components/permissions/prediction_service/permissions_aiv4_executor.cc`

### Common AI Model Components

- `//components/permissions/prediction_service/permissions_ai_encoder_base.h`
- `//components/permissions/prediction_service/permissions_ai_encoder_base.cc`