Snap for 11278173 from fba454f063a68c794214c6598181995591f0ebae to mainline-neuralnetworks-releaseaml_neu_341510000

Change-Id: I5ada25e7ca5d4de774368bd372152b1d6e1aefec

70 files changed, 1016 insertions, 906 deletions

diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java b/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
index 6af7c1ca..5ec7d5d5 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
@@ -166,4 +166,10 @@ public interface Flags {
     default Long getEncryptionKeyFetchPeriodSeconds() {
         return ENCRYPTION_KEY_FETCH_PERIOD_SECONDS;
     }
+
+    Boolean USE_BACKGROUND_ENCRYPTION_KEY_FETCH = true;
+
+    default Boolean getEnableBackgroundEncryptionKeyFetch() {
+        return USE_BACKGROUND_ENCRYPTION_KEY_FETCH;
+    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java b/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
index ade34909..43e0a937 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
@@ -38,6 +38,9 @@ public final class PhFlags implements Flags {
 
     static final String FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL =
             "fcp_encryption_key_download_url";
+
+    static final String ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH =
+            "enable_background_encryption_key_fetch";
     private static final PhFlags sSingleton = new PhFlags();
 
     /** Returns the singleton instance of the PhFlags. */
@@ -72,4 +75,13 @@ public final class PhFlags implements Flags {
                 /* defaultValue= */ ENCRYPTION_KEY_FETCH_URL
         );
     }
+
+    @Override
+    public Boolean getEnableBackgroundEncryptionKeyFetch() {
+        return DeviceConfig.getBoolean(
+                /* namespace= */ NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                /* name= */ ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH,
+                /* defaultValue= */ USE_BACKGROUND_ENCRYPTION_KEY_FETCH
+        );
+    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
index 69ca70a0..314392a3 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
@@ -151,6 +151,33 @@ public class FederatedTrainingTaskDao {
         }
     }
 
+    /** Delete a task from table based on population name and calling package. */
+    public FederatedTrainingTask findAndRemoveTaskByPopulationNameAndCallingPackage(
+            String populationName, String callingPackage) {
+        String selection =
+                FederatedTrainingTaskColumns.POPULATION_NAME
+                        + " = ? AND "
+                        + FederatedTrainingTaskColumns.APP_PACKAGE_NAME
+                        + " = ?";
+        String[] selectionArgs = {populationName, callingPackage};
+        FederatedTrainingTask task =
+                Iterables.getOnlyElement(getFederatedTrainingTask(selection, selectionArgs), null);
+        try {
+            if (task != null) {
+                deleteFederatedTrainingTask(selection, selectionArgs);
+            }
+            return task;
+        } catch (SQLException e) {
+            LogUtil.e(
+                    TAG,
+                    e,
+                    "Failed to delete federated training task by population name %s and ATP: %s",
+                    populationName,
+                    callingPackage);
+            return null;
+        }
+    }
+
     /** Delete a task from table based on population name and job scheduler id. */
     public FederatedTrainingTask findAndRemoveTaskByPopulationAndJobId(
             String populationName, int jobId) {
diff --git a/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java b/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java
index 6b8287ae..c48abb6f 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java
@@ -61,6 +61,10 @@ public class BackgroundKeyFetchJobService extends JobService {
 
     private final Injector mInjector;
 
+    public BackgroundKeyFetchJobService() {
+        mInjector = new Injector();
+    }
+
     @VisibleForTesting
     BackgroundKeyFetchJobService(Injector injector) {
         mInjector = injector;
@@ -137,6 +141,12 @@ public class BackgroundKeyFetchJobService extends JobService {
 
     /** Schedule the periodic background key fetch and delete job if it is not scheduled. */
     public static boolean scheduleJobIfNeeded(Context context, Flags flags) {
+        if (!flags.getEnableBackgroundEncryptionKeyFetch()) {
+            LogUtil.d(
+                    TAG,
+                    "Schedule encryption key job fetch is not enable in flags.");
+            return false;
+        }
         final JobScheduler jobScheduler = context.getSystemService(JobScheduler.class);
         if (jobScheduler == null) {
             LogUtil.e(TAG, "Failed to get job scheduler from system service.");
diff --git a/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java b/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
index 466a90db..476f54ff 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
@@ -152,8 +152,8 @@ public class FederatedComputeJobManager {
     public synchronized int onTrainerStartCalled(
             String callingPackageName, TrainingOptions trainingOptions) {
         FederatedTrainingTask existingTask =
-                mFederatedTrainingTaskDao.findAndRemoveTaskByPopulationName(
-                        trainingOptions.getPopulationName());
+                mFederatedTrainingTaskDao.findAndRemoveTaskByPopulationNameAndCallingPackage(
+                        trainingOptions.getPopulationName(), callingPackageName);
         Set<FederatedTrainingTask> trainingTasksToCancel = new HashSet<>();
         String populationName = trainingOptions.getPopulationName();
         long nowMs = mClock.currentTimeMillis();
@@ -286,11 +286,16 @@ public class FederatedComputeJobManager {
      */
     public synchronized int onTrainerStopCalled(String callingPackageName, String populationName) {
         FederatedTrainingTask taskToCancel =
-                mFederatedTrainingTaskDao.findAndRemoveTaskByPopulationName(populationName);
+                mFederatedTrainingTaskDao.findAndRemoveTaskByPopulationNameAndCallingPackage(
+                        populationName, callingPackageName);
         // If no matching task exists then there's nothing for us to do. This is not an error
         // case though.
         if (taskToCancel == null) {
-            LogUtil.i(TAG, "No matching task exists when cancel the job %s", populationName);
+            LogUtil.i(
+                    TAG,
+                    "No matching task exists when cancel the job (population: %s, ATP: %s",
+                    populationName,
+                    callingPackageName);
             return STATUS_SUCCESS;
         }
 
diff --git a/federatedcompute/src/com/android/federatedcompute/services/scheduling/JobSchedulerHelper.java b/federatedcompute/src/com/android/federatedcompute/services/scheduling/JobSchedulerHelper.java
index d009c910..08535790 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/scheduling/JobSchedulerHelper.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/scheduling/JobSchedulerHelper.java
@@ -39,7 +39,11 @@ public class JobSchedulerHelper {
     /** Schedules a task using JobScheduler. */
     public boolean scheduleTask(Context context, FederatedTrainingTask newTask) {
         JobInfo jobInfo = convertToJobInfo(context, newTask);
-        LogUtil.i(TAG, "Scheduling job %s", jobInfo.getId());
+        LogUtil.i(
+                TAG,
+                "Scheduling job %s (with minimum latency until next run %d milliseconds)",
+                jobInfo.getId(),
+                jobInfo.getMinLatencyMillis());
         return tryScheduleJob(context, jobInfo);
     }
 
diff --git a/federatedcompute/src/com/android/federatedcompute/services/security/KeyAttestation.java b/federatedcompute/src/com/android/federatedcompute/services/security/KeyAttestation.java
new file mode 100644
index 00000000..cdc43a31
--- /dev/null
+++ b/federatedcompute/src/com/android/federatedcompute/services/security/KeyAttestation.java
@@ -0,0 +1,188 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.federatedcompute.services.security;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.security.keystore.KeyGenParameterSpec;
+import android.security.keystore.KeyProperties;
+import android.security.keystore.StrongBoxUnavailableException;
+import android.util.Base64;
+
+import com.android.federatedcompute.internal.util.LogUtil;
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.io.IOException;
+import java.security.InvalidAlgorithmParameterException;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.NoSuchAlgorithmException;
+import java.security.NoSuchProviderException;
+import java.security.cert.Certificate;
+import java.security.cert.CertificateException;
+import java.security.spec.ECGenParameterSpec;
+import java.util.ArrayList;
+import java.util.List;
+
+public class KeyAttestation {
+
+    private static final String TAG = "KeyAttestation";
+    private static final String ANDROID_KEY_STORE = "AndroidKeyStore";
+
+    private static final String ODP_KEY_ALIAS = "ODPKeyAttestation";
+
+    private static volatile KeyAttestation sSingletonInstance;
+
+    private final boolean mUseStrongBox;
+
+    static class Injector {
+        KeyStore getKeyStore() throws KeyStoreException {
+            return KeyStore.getInstance(ANDROID_KEY_STORE);
+        }
+
+        KeyPairGenerator getKeyPairGenerator()
+                throws NoSuchAlgorithmException, NoSuchProviderException {
+            return KeyPairGenerator.getInstance(KeyProperties.KEY_ALGORITHM_EC, ANDROID_KEY_STORE);
+        }
+    }
+
+    private final Injector mInjector;
+
+    KeyAttestation(boolean useStrongBox, Injector injector) {
+        this.mUseStrongBox = useStrongBox;
+        this.mInjector = injector;
+    }
+
+    /**
+     * @return a singleton instance for KeyAttestation.
+     */
+    public static KeyAttestation getInstance(Context context) {
+        if (sSingletonInstance == null) {
+            synchronized (KeyAttestation.class) {
+                if (sSingletonInstance == null) {
+                    boolean useStrongBox = context.getPackageManager()
+                            .hasSystemFeature(PackageManager.FEATURE_STRONGBOX_KEYSTORE);
+                    sSingletonInstance = new KeyAttestation(useStrongBox, new Injector());
+                }
+            }
+        }
+        return sSingletonInstance;
+    }
+
+    /**
+     * For test only, return an singleton instance given the context and an injector for keyStore nd
+     * key pair generator.
+     */
+    @VisibleForTesting
+    static KeyAttestation getInstanceForTest(Context context, Injector injector) {
+        if (sSingletonInstance == null) {
+            synchronized (KeyAttestation.class) {
+                if (sSingletonInstance == null) {
+                    boolean useStrongBox = context.getPackageManager()
+                            .hasSystemFeature(PackageManager.FEATURE_STRONGBOX_KEYSTORE);
+                    return new KeyAttestation(useStrongBox, injector);
+                }
+            }
+        }
+        return sSingletonInstance;
+    }
+
+    /**
+     * Given a challenge, return a list of base64 encoded strings as the attestation record. The
+     * attestation is performed using a 256-bit Elliptical Curve (EC) key-pair generated by the
+     * secure keymaster.
+     */
+    public List<String> generateAttestationRecord(
+            final byte[] challenge, final String callingPackage) {
+        final String keyAlias = callingPackage + "-" + ODP_KEY_ALIAS;
+        KeyPair kp = generateHybridKey(challenge, keyAlias);
+        if (kp == null) {
+            return new ArrayList<>();
+        }
+        return getAttestationRecordFromKeyAlias(keyAlias);
+    }
+
+    @VisibleForTesting
+    KeyPair generateHybridKey(final byte[] challenge, final String keyAlias) {
+        try {
+            KeyPairGenerator keyPairGenerator = mInjector.getKeyPairGenerator();
+            keyPairGenerator.initialize(
+                    new KeyGenParameterSpec.Builder(
+                                    /* keystoreAlias= */ keyAlias,
+                                    /* purposes= */ KeyProperties.PURPOSE_SIGN)
+                            .setDigests(KeyProperties.DIGEST_SHA256)
+                            .setAlgorithmParameterSpec(new ECGenParameterSpec("secp256r1"))
+                            .setAttestationChallenge(challenge)
+                            // device properties are not specified when acquiring the challenge
+                            .setDevicePropertiesAttestationIncluded(false)
+                            .setIsStrongBoxBacked(mUseStrongBox)
+                            .build());
+            LogUtil.e(TAG, keyPairGenerator.getAlgorithm());
+            return keyPairGenerator.generateKeyPair();
+        } catch (NoSuchAlgorithmException
+                | NoSuchProviderException
+                | InvalidAlgorithmParameterException e) {
+            LogUtil.e(TAG, "Failed to generate EC key for attestation: " + e.getMessage());
+        } catch (StrongBoxUnavailableException e) {
+            LogUtil.e(
+                    TAG,
+                    "Strong box not available on device but isStrongBox is set to true: "
+                            + e.getMessage());
+        }
+        return null;
+    }
+
+    @VisibleForTesting
+    List<String> getAttestationRecordFromKeyAlias(String keyAlias) {
+        try {
+            KeyStore keyStore = mInjector.getKeyStore();
+            keyStore.load(null);
+            Certificate[] certificateChain = keyStore.getCertificateChain(keyAlias);
+            if (certificateChain == null) {
+                return new ArrayList<>();
+            }
+            ArrayList<String> attestationRecord = new ArrayList<>();
+            for (Certificate certificate : certificateChain) {
+                attestationRecord.add(
+                        Base64.encodeToString(certificate.getEncoded(), Base64.NO_WRAP));
+            }
+            return attestationRecord;
+        } catch (CertificateException e) {
+            LogUtil.e(
+                    TAG,
+                    "CertificateException when"
+                            + "generate certs for attestation: " + e.getMessage());
+        } catch (IOException e) {
+            LogUtil.e(
+                    TAG, "IOException when "
+                            + "generate certs for attestation: " + e.getMessage());
+        } catch (KeyStoreException e) {
+            LogUtil.e(
+                    TAG,
+                    "KeystoreException when "
+                            + "generate certs for attestation: " + e.getMessage());
+        } catch (NoSuchAlgorithmException e) {
+            LogUtil.e(
+                    TAG,
+                    "NoSuchAlgorithmException when"
+                            + "generate certs for attestation: " + e.getMessage());
+        }
+        return new ArrayList<>();
+    }
+}
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
index 4244aff4..a6583bab 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
@@ -244,6 +244,8 @@ public class FederatedComputeWorker {
             TrainingRun run, CheckinResult checkinResult) {
         // Stop processing if have rejection Info
         if (checkinResult.getRejectionInfo() != null) {
+            LogUtil.d(TAG, "job %d was rejected during check in, reason %s",
+                    run.mTask.jobId(), checkinResult.getRejectionInfo().getReason());
             mJobManager.onTrainingCompleted(
                     run.mTask.jobId(),
                     run.mTask.populationName(),
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
index 79966d99..140a523a 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
@@ -50,7 +50,9 @@ import java.util.ArrayList;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
 
-/** The implementation of {@link IsolatedTrainingService}. */
+/**
+ * The implementation of {@link IsolatedTrainingService}.
+ */
 public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
     private static final String TAG = IsolatedTrainingServiceImpl.class.getSimpleName();
     private final AtomicBoolean mInterruptFlag = new AtomicBoolean(false);
@@ -143,12 +145,18 @@ public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
                         for (ExampleConsumption exampleConsumption : exampleConsumptionArrayList) {
                             numExamples += exampleConsumption.getExampleCount();
                         }
-                        LogUtil.i(
-                                TAG,
-                                "training task %s: result %s, used %d examples",
-                                populationName,
-                                result.toString(),
-                                numExamples);
+                        if (result.getContributionResult() == ContributionResult.SUCCESS) {
+                            LogUtil.i(
+                                    TAG,
+                                    "training task %s: result %s, used %d examples",
+                                    populationName,
+                                    result.getContributionResult(),
+                                    numExamples);
+                        } else {
+                            LogUtil.i(TAG, "training task %s: result %s, error message %s",
+                                    populationName, result.getContributionResult(),
+                                    result.getErrorMessage());
+                        }
                         bundle.putParcelableArrayList(
                                 ClientConstants.EXTRA_EXAMPLE_CONSUMPTION_LIST,
                                 exampleConsumptionArrayList);
diff --git a/framework/api/current.txt b/framework/api/current.txt
index cee99647..d802177e 100644
--- a/framework/api/current.txt
+++ b/framework/api/current.txt
@@ -1,251 +1 @@
 // Signature format: 2.0
-package android.adservices.ondevicepersonalization {
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class AppInfo implements android.os.Parcelable {
-    method public int describeContents();
-    method @NonNull public boolean isInstalled();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.AppInfo> CREATOR;
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class DownloadCompletedInput {
-    method @NonNull public java.util.Map<java.lang.String,byte[]> getData();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class DownloadCompletedOutput {
-    method @NonNull public java.util.List<java.lang.String> getRetainedKeys();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class DownloadCompletedOutput.Builder {
-    ctor public DownloadCompletedOutput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.DownloadCompletedOutput.Builder addRetainedKey(@NonNull String);
-    method @NonNull public android.adservices.ondevicepersonalization.DownloadCompletedOutput build();
-    method @NonNull public android.adservices.ondevicepersonalization.DownloadCompletedOutput.Builder setRetainedKeys(@NonNull java.util.List<java.lang.String>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class EventInput {
-    method @NonNull public android.os.PersistableBundle getParameters();
-    method @Nullable public android.adservices.ondevicepersonalization.RequestLogRecord getRequestLogRecord();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class EventLogRecord implements android.os.Parcelable {
-    method public int describeContents();
-    method @Nullable public android.content.ContentValues getData();
-    method @Nullable public android.adservices.ondevicepersonalization.RequestLogRecord getRequestLogRecord();
-    method @IntRange(from=0) public int getRowIndex();
-    method @NonNull public java.time.Instant getTime();
-    method @IntRange(from=1, to=127) public int getType();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.EventLogRecord> CREATOR;
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class EventLogRecord.Builder {
-    ctor public EventLogRecord.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.EventLogRecord build();
-    method @NonNull public android.adservices.ondevicepersonalization.EventLogRecord.Builder setData(@NonNull android.content.ContentValues);
-    method @NonNull public android.adservices.ondevicepersonalization.EventLogRecord.Builder setRequestLogRecord(@NonNull android.adservices.ondevicepersonalization.RequestLogRecord);
-    method @NonNull public android.adservices.ondevicepersonalization.EventLogRecord.Builder setRowIndex(@IntRange(from=0) int);
-    method @NonNull public android.adservices.ondevicepersonalization.EventLogRecord.Builder setType(@IntRange(from=1, to=127) int);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class EventOutput {
-    method @Nullable public android.adservices.ondevicepersonalization.EventLogRecord getEventLogRecord();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class EventOutput.Builder {
-    ctor public EventOutput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.EventOutput build();
-    method @NonNull public android.adservices.ondevicepersonalization.EventOutput.Builder setEventLogRecord(@NonNull android.adservices.ondevicepersonalization.EventLogRecord);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class EventUrlProvider {
-    method @NonNull @WorkerThread public android.net.Uri createEventTrackingUrlWithRedirect(@NonNull android.os.PersistableBundle, @Nullable android.net.Uri);
-    method @NonNull @WorkerThread public android.net.Uri createEventTrackingUrlWithResponse(@NonNull android.os.PersistableBundle, @Nullable byte[], @Nullable String);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class ExecuteInput {
-    method @NonNull public String getAppPackageName();
-    method @NonNull public android.os.PersistableBundle getAppParams();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class ExecuteOutput {
-    method @NonNull public java.util.List<android.adservices.ondevicepersonalization.RenderingConfig> getRenderingConfigs();
-    method @Nullable public android.adservices.ondevicepersonalization.RequestLogRecord getRequestLogRecord();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class ExecuteOutput.Builder {
-    ctor public ExecuteOutput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.ExecuteOutput.Builder addRenderingConfig(@NonNull android.adservices.ondevicepersonalization.RenderingConfig);
-    method @NonNull public android.adservices.ondevicepersonalization.ExecuteOutput build();
-    method @NonNull public android.adservices.ondevicepersonalization.ExecuteOutput.Builder setRenderingConfigs(@NonNull java.util.List<android.adservices.ondevicepersonalization.RenderingConfig>);
-    method @NonNull public android.adservices.ondevicepersonalization.ExecuteOutput.Builder setRequestLogRecord(@NonNull android.adservices.ondevicepersonalization.RequestLogRecord);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class FederatedComputeInput {
-    method @NonNull public String getPopulationName();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class FederatedComputeInput.Builder {
-    ctor public FederatedComputeInput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.FederatedComputeInput build();
-    method @NonNull public android.adservices.ondevicepersonalization.FederatedComputeInput.Builder setPopulationName(@NonNull String);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class FederatedComputeScheduler {
-    method @WorkerThread public void cancel(@NonNull String);
-    method @WorkerThread public void schedule(@NonNull android.adservices.ondevicepersonalization.FederatedComputeScheduler.Params, @NonNull android.adservices.ondevicepersonalization.FederatedComputeInput);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static class FederatedComputeScheduler.Params {
-    ctor public FederatedComputeScheduler.Params(@NonNull android.adservices.ondevicepersonalization.TrainingInterval);
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingInterval getTrainingInterval();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public abstract class IsolatedService extends android.app.Service {
-    ctor public IsolatedService();
-    method @NonNull public final android.adservices.ondevicepersonalization.EventUrlProvider getEventUrlProvider(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @NonNull public final android.adservices.ondevicepersonalization.FederatedComputeScheduler getFederatedComputeScheduler(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @NonNull public final android.adservices.ondevicepersonalization.MutableKeyValueStore getLocalData(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @NonNull public final android.adservices.ondevicepersonalization.LogReader getLogReader(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @NonNull public final android.adservices.ondevicepersonalization.KeyValueStore getRemoteData(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @Nullable public final android.adservices.ondevicepersonalization.UserData getUserData(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-    method @Nullable public android.os.IBinder onBind(@NonNull android.content.Intent);
-    method @NonNull public abstract android.adservices.ondevicepersonalization.IsolatedWorker onRequest(@NonNull android.adservices.ondevicepersonalization.RequestToken);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public interface IsolatedWorker {
-    method public default void onDownloadCompleted(@NonNull android.adservices.ondevicepersonalization.DownloadCompletedInput, @NonNull java.util.function.Consumer<android.adservices.ondevicepersonalization.DownloadCompletedOutput>);
-    method public default void onEvent(@NonNull android.adservices.ondevicepersonalization.EventInput, @NonNull java.util.function.Consumer<android.adservices.ondevicepersonalization.EventOutput>);
-    method public default void onExecute(@NonNull android.adservices.ondevicepersonalization.ExecuteInput, @NonNull java.util.function.Consumer<android.adservices.ondevicepersonalization.ExecuteOutput>);
-    method public default void onRender(@NonNull android.adservices.ondevicepersonalization.RenderInput, @NonNull java.util.function.Consumer<android.adservices.ondevicepersonalization.RenderOutput>);
-    method public default void onTrainingExamples(@NonNull android.adservices.ondevicepersonalization.TrainingExamplesInput, @NonNull java.util.function.Consumer<android.adservices.ondevicepersonalization.TrainingExamplesOutput>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public interface KeyValueStore {
-    method @Nullable @WorkerThread public byte[] get(@NonNull String);
-    method @NonNull @WorkerThread public java.util.Set<java.lang.String> keySet();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class LogReader {
-    method @NonNull @WorkerThread public java.util.List<android.adservices.ondevicepersonalization.EventLogRecord> getJoinedEvents(@NonNull java.time.Instant, @NonNull java.time.Instant);
-    method @NonNull @WorkerThread public java.util.List<android.adservices.ondevicepersonalization.RequestLogRecord> getRequests(@NonNull java.time.Instant, @NonNull java.time.Instant);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public interface MutableKeyValueStore extends android.adservices.ondevicepersonalization.KeyValueStore {
-    method @Nullable @WorkerThread public byte[] put(@NonNull String, @NonNull byte[]);
-    method @Nullable @WorkerThread public byte[] remove(@NonNull String);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class OnDevicePersonalizationException extends java.lang.Exception {
-    method public int getErrorCode();
-    field public static final int ERROR_ISOLATED_SERVICE_FAILED = 1; // 0x1
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class OnDevicePersonalizationManager {
-    method public void execute(@NonNull android.content.ComponentName, @NonNull android.os.PersistableBundle, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<java.util.List<android.adservices.ondevicepersonalization.SurfacePackageToken>,java.lang.Exception>);
-    method public void requestSurfacePackage(@NonNull android.adservices.ondevicepersonalization.SurfacePackageToken, @NonNull android.os.IBinder, int, int, int, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<android.view.SurfaceControlViewHost.SurfacePackage,java.lang.Exception>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class RenderInput {
-    method public int getHeight();
-    method @Nullable public android.adservices.ondevicepersonalization.RenderingConfig getRenderingConfig();
-    method public int getRenderingConfigIndex();
-    method public int getWidth();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class RenderOutput {
-    method @Nullable public String getContent();
-    method @Nullable public String getTemplateId();
-    method @NonNull public android.os.PersistableBundle getTemplateParams();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class RenderOutput.Builder {
-    ctor public RenderOutput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.RenderOutput build();
-    method @NonNull public android.adservices.ondevicepersonalization.RenderOutput.Builder setContent(@NonNull String);
-    method @NonNull public android.adservices.ondevicepersonalization.RenderOutput.Builder setTemplateId(@NonNull String);
-    method @NonNull public android.adservices.ondevicepersonalization.RenderOutput.Builder setTemplateParams(@NonNull android.os.PersistableBundle);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class RenderingConfig implements android.os.Parcelable {
-    method public int describeContents();
-    method @NonNull public java.util.List<java.lang.String> getKeys();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.RenderingConfig> CREATOR;
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class RenderingConfig.Builder {
-    ctor public RenderingConfig.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.RenderingConfig.Builder addKey(@NonNull String);
-    method @NonNull public android.adservices.ondevicepersonalization.RenderingConfig build();
-    method @NonNull public android.adservices.ondevicepersonalization.RenderingConfig.Builder setKeys(@NonNull java.util.List<java.lang.String>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class RequestLogRecord implements android.os.Parcelable {
-    method public int describeContents();
-    method @NonNull public java.util.List<android.content.ContentValues> getRows();
-    method @NonNull public java.time.Instant getTime();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.RequestLogRecord> CREATOR;
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class RequestLogRecord.Builder {
-    ctor public RequestLogRecord.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.RequestLogRecord.Builder addRow(@NonNull android.content.ContentValues);
-    method @NonNull public android.adservices.ondevicepersonalization.RequestLogRecord build();
-    method @NonNull public android.adservices.ondevicepersonalization.RequestLogRecord.Builder setRows(@NonNull java.util.List<android.content.ContentValues>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class RequestToken {
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class SurfacePackageToken {
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class TrainingExamplesInput {
-    method @NonNull public String getPopulationName();
-    method @Nullable public byte[] getResumptionToken();
-    method @NonNull public String getTaskName();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class TrainingExamplesOutput {
-    method @NonNull public java.util.List<byte[]> getResumptionTokens();
-    method @NonNull public java.util.List<byte[]> getTrainingExamples();
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class TrainingExamplesOutput.Builder {
-    ctor public TrainingExamplesOutput.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingExamplesOutput.Builder addResumptionToken(@NonNull byte[]);
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingExamplesOutput.Builder addTrainingExample(@NonNull byte[]);
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingExamplesOutput build();
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingExamplesOutput.Builder setResumptionTokens(@NonNull java.util.List<byte[]>);
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingExamplesOutput.Builder setTrainingExamples(@NonNull java.util.List<byte[]>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class TrainingInterval {
-    method @NonNull public java.time.Duration getMinimumInterval();
-    method public int getSchedulingMode();
-    field public static final int SCHEDULING_MODE_ONE_TIME = 1; // 0x1
-    field public static final int SCHEDULING_MODE_RECURRENT = 2; // 0x2
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public static final class TrainingInterval.Builder {
-    ctor public TrainingInterval.Builder();
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingInterval build();
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingInterval.Builder setMinimumInterval(@NonNull java.time.Duration);
-    method @NonNull public android.adservices.ondevicepersonalization.TrainingInterval.Builder setSchedulingMode(int);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public final class UserData implements android.os.Parcelable {
-    method public int describeContents();
-    method @NonNull public java.util.Map<java.lang.String,android.adservices.ondevicepersonalization.AppInfo> getAppInfos();
-    method @IntRange(from=0) public long getAvailableStorageBytes();
-    method @IntRange(from=0, to=100) public int getBatteryPercentage();
-    method @NonNull public String getCarrier();
-    method public int getOrientation();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.UserData> CREATOR;
-  }
-
-}
-
diff --git a/framework/api/system-current.txt b/framework/api/system-current.txt
index e7050287..d802177e 100644
--- a/framework/api/system-current.txt
+++ b/framework/api/system-current.txt
@@ -1,13 +1 @@
 // Signature format: 2.0
-package android.adservices.ondevicepersonalization {
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class OnDevicePersonalizationConfigManager {
-    method @FlaggedApi("enable_ondevicepersonalization_apis") @RequiresPermission(android.adservices.ondevicepersonalization.OnDevicePersonalizationPermissions.MODIFY_ONDEVICEPERSONALIZATION_STATE) public void setPersonalizationEnabled(boolean, @NonNull java.util.concurrent.Executor, @NonNull android.os.OutcomeReceiver<java.lang.Void,java.lang.Exception>);
-  }
-
-  @FlaggedApi("enable_ondevicepersonalization_apis") public class OnDevicePersonalizationPermissions {
-    field @FlaggedApi("enable_ondevicepersonalization_apis") public static final String MODIFY_ONDEVICEPERSONALIZATION_STATE = "android.permission.ondevicepersonalization.MODIFY_ONDEVICEPERSONALIZATION_STATE";
-  }
-
-}
-
diff --git a/framework/java/android/adservices/ondevicepersonalization/AppInfo.java b/framework/java/android/adservices/ondevicepersonalization/AppInfo.java
index 4a628d8d..d5648386 100644
--- a/framework/java/android/adservices/ondevicepersonalization/AppInfo.java
+++ b/framework/java/android/adservices/ondevicepersonalization/AppInfo.java
@@ -28,6 +28,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 /**
  * Information about apps.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedInput.java b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedInput.java
index 14dfb320..f413dc97 100644
--- a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedInput.java
@@ -32,6 +32,7 @@ import java.util.Map;
  * The input data for {@link
  * IsolatedWorker#onDownloadCompleted(DownloadCompletedInput, java.util.function.Consumer)}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
index 5c503fcb..c3a26e20 100644
--- a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
@@ -31,6 +31,7 @@ import java.util.List;
  * The result returned by {@link
  * IsolatedWorker#onDownloadCompleted(DownloadCompletedInput, java.util.function.Consumer)}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventInput.java b/framework/java/android/adservices/ondevicepersonalization/EventInput.java
index 33057ad6..1fd6e448 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventInput.java
@@ -29,6 +29,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 /**
  * The input data for {@link
  * IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}.
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java b/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
index bb1c6101..3735d669 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
@@ -42,6 +42,8 @@ import java.time.Instant;
  * The contents of the EVENTS table can be
  * consumed by Federated Learning facilitated model training, or Federated Analytics facilitated
  * cross-device statistical analysis.
+ *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventOutput.java b/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
index 9fe978d6..e3302337 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
@@ -25,6 +25,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 
 /**
  *  The result returned by {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}.
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventUrlProvider.java b/framework/java/android/adservices/ondevicepersonalization/EventUrlProvider.java
index 930c434f..7605c2c3 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventUrlProvider.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventUrlProvider.java
@@ -40,6 +40,7 @@ import java.util.concurrent.BlockingQueue;
  * {@code IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}, and log the returned
  * output in the EVENTS table.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class EventUrlProvider {
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
index 825b25a2..e939a6b9 100644
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
@@ -28,6 +28,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 /**
  * The input data for {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
index b6ee67dd..e5f0f3f5 100644
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
@@ -34,6 +34,7 @@ import java.util.List;
  * {@code OnDevicePersonalizationManager#execute(ComponentName, PersistableBundle,
  * java.util.concurrent.Executor, OutcomeReceiver)}
  * from a client app.
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
index 6057b2c1..d91892be 100644
--- a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
@@ -24,7 +24,10 @@ import android.annotation.NonNull;
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
-/** The input data for {@link FederatedComputeScheduler#schedule}. */
+/**
+ * The input data for {@link FederatedComputeScheduler#schedule}.
+ * @hide
+ */
 @DataClass(genBuilder = true, genEqualsHashCode = true)
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public final class FederatedComputeInput {
diff --git a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
index 6e6f2bef..4ed043da 100644
--- a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
+++ b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
@@ -33,6 +33,7 @@ import java.util.concurrent.CountDownLatch;
 /**
  * Handles scheduling federated compute jobs. See {@link
  * IsolatedService#getFederatedComputeScheduler}.
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class FederatedComputeScheduler {
diff --git a/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java b/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
index 5a8fd11a..832adf4a 100644
--- a/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
+++ b/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
@@ -55,6 +55,8 @@ import java.util.function.Function;
  * by Federated Learning for model training.
  * Client apps use {@link OnDevicePersonalizationManager} to interact with an {@link
  * IsolatedService}.
+ *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public abstract class IsolatedService extends Service {
diff --git a/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java b/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
index 11a651eb..90a27675 100644
--- a/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
+++ b/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
@@ -30,6 +30,7 @@ import java.util.function.Consumer;
  * IsolatedService} calls the method on a Binder thread and the {@link IsolatedWorker} should
  * offload long running operations to a worker thread. The consumer parameter of each method is used
  * to return results.
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public interface IsolatedWorker {
diff --git a/framework/java/android/adservices/ondevicepersonalization/KeyValueStore.java b/framework/java/android/adservices/ondevicepersonalization/KeyValueStore.java
index b50f70e2..e943194e 100644
--- a/framework/java/android/adservices/ondevicepersonalization/KeyValueStore.java
+++ b/framework/java/android/adservices/ondevicepersonalization/KeyValueStore.java
@@ -32,6 +32,7 @@ import java.util.Set;
  *
  * @see IsolatedService#getRemoteData(RequestToken)
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public interface KeyValueStore {
diff --git a/framework/java/android/adservices/ondevicepersonalization/LogReader.java b/framework/java/android/adservices/ondevicepersonalization/LogReader.java
index bf25efcf..c06c48d2 100644
--- a/framework/java/android/adservices/ondevicepersonalization/LogReader.java
+++ b/framework/java/android/adservices/ondevicepersonalization/LogReader.java
@@ -43,6 +43,7 @@ import java.util.concurrent.BlockingQueue;
  *
  * @see IsolatedService#getLogReader(RequestToken)
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class LogReader {
diff --git a/framework/java/android/adservices/ondevicepersonalization/MutableKeyValueStore.java b/framework/java/android/adservices/ondevicepersonalization/MutableKeyValueStore.java
index e63bcc6f..b67df7e8 100644
--- a/framework/java/android/adservices/ondevicepersonalization/MutableKeyValueStore.java
+++ b/framework/java/android/adservices/ondevicepersonalization/MutableKeyValueStore.java
@@ -30,6 +30,7 @@ import android.annotation.WorkerThread;
  *
  * @see IsolatedService#getLocalData(RequestToken)
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public interface MutableKeyValueStore extends KeyValueStore {
diff --git a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationConfigManager.java b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationConfigManager.java
index f6442fec..82a38ce6 100644
--- a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationConfigManager.java
+++ b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationConfigManager.java
@@ -26,7 +26,6 @@ import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
-import android.annotation.SystemApi;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -51,7 +50,6 @@ import java.util.concurrent.TimeUnit;
  *
  * @hide
  */
-@SystemApi
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class OnDevicePersonalizationConfigManager {
     /** @hide */
diff --git a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationException.java b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationException.java
index e54f76bf..c88fa381 100644
--- a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationException.java
+++ b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationException.java
@@ -27,6 +27,7 @@ import java.lang.annotation.RetentionPolicy;
 /**
  * Exception thrown by OnDevicePersonalization APIs.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class OnDevicePersonalizationException extends Exception {
diff --git a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
index ab772f3a..2a267a92 100644
--- a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
+++ b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
@@ -52,6 +52,8 @@ import java.util.concurrent.Executor;
  * persistent results to on-device storage which can be consumed by Federated Analytics for
  * cross-device statistical analysis or by Federated Learning for model training. The displayed
  * content and the persistent output are both not directly accessible by the calling app.
+ *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class OnDevicePersonalizationManager {
diff --git a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationPermissions.java b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationPermissions.java
index 694e336c..b87e0736 100644
--- a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationPermissions.java
+++ b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationPermissions.java
@@ -20,7 +20,6 @@ import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ON
 
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
-import android.annotation.SystemApi;
 import android.content.Context;
 import android.content.pm.PackageManager;
 
@@ -29,7 +28,6 @@ import android.content.pm.PackageManager;
  *
  * @hide
 */
-@SystemApi
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class OnDevicePersonalizationPermissions {
     private OnDevicePersonalizationPermissions() {}
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderInput.java b/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
index 0c5824de..58f0e8eb 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
@@ -28,6 +28,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
  * The input data for
  * {@link IsolatedWorker#onRender(RenderInput, java.util.function.Consumer)}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java b/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
index 33c8c8b3..f1ab9e13 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
@@ -30,6 +30,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
  * The result returned by
  * {@link IsolatedWorker#onRender(RenderInput, java.util.function.Consumer)}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderingConfig.java b/framework/java/android/adservices/ondevicepersonalization/RenderingConfig.java
index ae54c77a..a79cbcaf 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RenderingConfig.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderingConfig.java
@@ -35,6 +35,7 @@ import java.util.List;
  * {@link IsolatedWorker#onRender(RenderInput, java.util.function.Consumer)} to identify the
  * content to be displayed in a single {@link android.view.View}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java b/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
index 618462e8..bab8afaa 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
@@ -39,6 +39,7 @@ import java.util.List;
  * the REQUESTS table can be consumed by Federated Learning facilitated model training,
  * or Federated Analytics facilitated cross-device statistical analysis.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
diff --git a/framework/java/android/adservices/ondevicepersonalization/RequestToken.java b/framework/java/android/adservices/ondevicepersonalization/RequestToken.java
index d3e23310..04c67f6d 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RequestToken.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RequestToken.java
@@ -30,6 +30,8 @@ import java.util.Objects;
 /**
  * An opaque token that identifies the current request to an {@link IsolatedService}. This token
  * must be passed as a parameter to all service methods that depend on per-request state.
+ *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class RequestToken {
diff --git a/framework/java/android/adservices/ondevicepersonalization/SurfacePackageToken.java b/framework/java/android/adservices/ondevicepersonalization/SurfacePackageToken.java
index 48406c3d..c8340b9c 100644
--- a/framework/java/android/adservices/ondevicepersonalization/SurfacePackageToken.java
+++ b/framework/java/android/adservices/ondevicepersonalization/SurfacePackageToken.java
@@ -25,6 +25,7 @@ import android.annotation.NonNull;
  * An opaque reference to content that can be displayed in a {@link android.view.SurfaceView}. This
  * maps to a {@link RenderingConfig} returned by an {@link IsolatedService}.
  *
+ * @hide
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class SurfacePackageToken {
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java
index bd8ca99e..6b6dd257 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java
@@ -25,7 +25,11 @@ import android.annotation.Nullable;
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
-/** The input data for {@link IsolatedWorker#onTrainingExamples}. */
+/**
+ * The input data for {@link IsolatedWorker#onTrainingExamples}.
+ *
+ * @hide
+ */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
 public final class TrainingExamplesInput {
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java
index deea25cb..246aabab 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java
@@ -28,7 +28,10 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 import java.util.Collections;
 import java.util.List;
 
-/** The output data of {@link IsolatedWorker#onTrainingExamples} */
+/**
+ * The output data of {@link IsolatedWorker#onTrainingExamples}
+ * @hide
+ */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
 public final class TrainingExamplesOutput {
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java b/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
index 890620e7..2f822c76 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
@@ -26,7 +26,11 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 
 import java.time.Duration;
 
-/** Training interval settings required for federated computation jobs. */
+/**
+ * Training interval settings required for federated computation jobs.
+ *
+ * @hide
+ */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genHiddenConstDefs = true, genEqualsHashCode = true)
 public final class TrainingInterval {
diff --git a/framework/java/android/adservices/ondevicepersonalization/UserData.java b/framework/java/android/adservices/ondevicepersonalization/UserData.java
index 644ea74a..442c279a 100644
--- a/framework/java/android/adservices/ondevicepersonalization/UserData.java
+++ b/framework/java/android/adservices/ondevicepersonalization/UserData.java
@@ -43,7 +43,7 @@ import java.util.Map;
 
 /**
  * User data provided by the platform to an {@link IsolatedService}.
- *
+ * @hide
  */
 // This class should be updated with the Kotlin mirror
 // {@link com.android.ondevicepersonalization.services.policyengine.data.UserData}.
diff --git a/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java b/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
index d74f0c71..7858229c 100644
--- a/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
+++ b/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
@@ -46,6 +46,7 @@ import com.android.ondevicepersonalization.services.manifest.AppManifestConfigHe
 import com.android.ondevicepersonalization.services.policyengine.UserDataAccessor;
 import com.android.ondevicepersonalization.services.process.IsolatedServiceInfo;
 import com.android.ondevicepersonalization.services.process.ProcessRunner;
+import com.android.ondevicepersonalization.services.process.ProcessRunnerImpl;
 import com.android.ondevicepersonalization.services.statsd.ApiCallStats;
 import com.android.ondevicepersonalization.services.statsd.OdpStatsdLogger;
 import com.android.ondevicepersonalization.services.util.Clock;
@@ -98,7 +99,7 @@ class OdpWebViewClient extends WebViewClient {
         }
 
         ProcessRunner getProcessRunner() {
-            return ProcessRunner.getInstance();
+            return ProcessRunnerImpl.getInstance();
         }
     }
 
diff --git a/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java b/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
index b05527e4..0c45f960 100644
--- a/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
+++ b/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
@@ -41,6 +41,7 @@ import com.android.ondevicepersonalization.services.manifest.AppManifestConfigHe
 import com.android.ondevicepersonalization.services.policyengine.UserDataAccessor;
 import com.android.ondevicepersonalization.services.process.IsolatedServiceInfo;
 import com.android.ondevicepersonalization.services.process.ProcessRunner;
+import com.android.ondevicepersonalization.services.process.ProcessRunnerImpl;
 import com.android.ondevicepersonalization.services.statsd.ApiCallStats;
 import com.android.ondevicepersonalization.services.statsd.OdpStatsdLogger;
 import com.android.ondevicepersonalization.services.util.Clock;
@@ -87,7 +88,7 @@ public class OnDevicePersonalizationDataProcessingAsyncCallable implements Async
         }
 
         ProcessRunner getProcessRunner() {
-            return ProcessRunner.getInstance();
+            return ProcessRunnerImpl.getInstance();
         }
     }
 
@@ -248,8 +249,10 @@ public class OnDevicePersonalizationDataProcessingAsyncCallable implements Async
                 filteredList.add(vendorDataMap.get(key));
             }
         }
-        mDao.batchUpdateOrInsertVendorDataTransaction(filteredList, retainedKeys,
-                syncToken);
+        boolean transactionResult = mDao.batchUpdateOrInsertVendorDataTransaction(
+                filteredList, retainedKeys, syncToken);
+        sLogger.d(TAG + ": filter and store data completed, transaction successful: "
+                + transactionResult);
         return null;
     }
 
diff --git a/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java b/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
index 5d63c651..fed0a940 100644
--- a/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
+++ b/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
@@ -40,6 +40,7 @@ import com.android.ondevicepersonalization.services.manifest.AppManifestConfigHe
 import com.android.ondevicepersonalization.services.policyengine.UserDataAccessor;
 import com.android.ondevicepersonalization.services.process.IsolatedServiceInfo;
 import com.android.ondevicepersonalization.services.process.ProcessRunner;
+import com.android.ondevicepersonalization.services.process.ProcessRunnerImpl;
 import com.android.ondevicepersonalization.services.statsd.ApiCallStats;
 import com.android.ondevicepersonalization.services.statsd.OdpStatsdLogger;
 import com.android.ondevicepersonalization.services.util.Clock;
@@ -77,7 +78,7 @@ public final class OdpExampleStoreService extends ExampleStoreService {
         }
 
         ProcessRunner getProcessRunner() {
-            return ProcessRunner.getInstance();
+            return ProcessRunnerImpl.getInstance();
         }
     }
 
diff --git a/src/com/android/ondevicepersonalization/services/process/OnDevicePersonalizationPlugin.java b/src/com/android/ondevicepersonalization/services/process/OnDevicePersonalizationPlugin.java
index 440fd832..c1a07d95 100644
--- a/src/com/android/ondevicepersonalization/services/process/OnDevicePersonalizationPlugin.java
+++ b/src/com/android/ondevicepersonalization/services/process/OnDevicePersonalizationPlugin.java
@@ -56,21 +56,21 @@ public class OnDevicePersonalizationPlugin implements Plugin {
         mPluginContext = pluginContext;
 
         try {
-            String className = input.getString(ProcessRunner.PARAM_CLASS_NAME_KEY);
+            String className = input.getString(ProcessRunnerImpl.PARAM_CLASS_NAME_KEY);
             if (className == null || className.isEmpty()) {
                 sLogger.e(TAG + ": className missing.");
                 sendErrorResult(FailureType.ERROR_EXECUTING_PLUGIN);
                 return;
             }
 
-            int operation = input.getInt(ProcessRunner.PARAM_OPERATION_KEY);
+            int operation = input.getInt(ProcessRunnerImpl.PARAM_OPERATION_KEY);
             if (operation == 0) {
                 sLogger.e(TAG + ": operation missing or invalid.");
                 sendErrorResult(FailureType.ERROR_EXECUTING_PLUGIN);
                 return;
             }
 
-            Bundle serviceParams = input.getParcelable(ProcessRunner.PARAM_SERVICE_INPUT,
+            Bundle serviceParams = input.getParcelable(ProcessRunnerImpl.PARAM_SERVICE_INPUT,
                     Bundle.class);
             if (serviceParams == null) {
                 sLogger.e(TAG + ": Missing service input.");
diff --git a/src/com/android/ondevicepersonalization/services/process/ProcessRunner.java b/src/com/android/ondevicepersonalization/services/process/ProcessRunner.java
index 8de09e94..26f5128a 100644
--- a/src/com/android/ondevicepersonalization/services/process/ProcessRunner.java
+++ b/src/com/android/ondevicepersonalization/services/process/ProcessRunner.java
@@ -16,222 +16,24 @@
 
 package com.android.ondevicepersonalization.services.process;
 
-import android.adservices.ondevicepersonalization.Constants;
 import android.annotation.NonNull;
-import android.annotation.Nullable;
-import android.content.Context;
 import android.os.Bundle;
 
-
-import androidx.concurrent.futures.CallbackToFutureAdapter;
-
-import com.android.ondevicepersonalization.internal.util.LoggerFactory;
-import com.android.ondevicepersonalization.libraries.plugin.FailureType;
-import com.android.ondevicepersonalization.libraries.plugin.PluginCallback;
-import com.android.ondevicepersonalization.libraries.plugin.PluginController;
-import com.android.ondevicepersonalization.libraries.plugin.PluginInfo;
-import com.android.ondevicepersonalization.libraries.plugin.PluginManager;
-import com.android.ondevicepersonalization.libraries.plugin.impl.PluginManagerImpl;
-import com.android.ondevicepersonalization.services.OdpServiceException;
-import com.android.ondevicepersonalization.services.OnDevicePersonalizationApplication;
-import com.android.ondevicepersonalization.services.util.Clock;
-import com.android.ondevicepersonalization.services.util.MonotonicClock;
-
-import com.google.common.collect.ImmutableList;
-import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 
-import java.util.Objects;
-
-/** Utilities to support loading and executing plugins. */
-public class ProcessRunner {
-    private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
-    private static final String TAG = "ProcessUtils";
-    private static final String ENTRY_POINT_CLASS =
-            "com.android.ondevicepersonalization.services.process.OnDevicePersonalizationPlugin";
-
-    public static final String PARAM_CLASS_NAME_KEY = "param.classname";
-    public static final String PARAM_OPERATION_KEY = "param.operation";
-    public static final String PARAM_SERVICE_INPUT = "param.service_input";
-
-    @NonNull private Context mApplicationContext;
-
-    private static volatile ProcessRunner sProcessRunner;
-    private static volatile PluginManager sPluginManager;
-
-    static class Injector {
-        Clock getClock() {
-            return MonotonicClock.getInstance();
-        }
-    }
-
-    private final Injector mInjector;
-
-    /** Creates a ProcessRunner. */
-    ProcessRunner(
-            @NonNull Context applicationContext,
-            @NonNull Injector injector) {
-        mApplicationContext = Objects.requireNonNull(applicationContext);
-        mInjector = Objects.requireNonNull(injector);
-    }
-
-    /** Returns the global ProcessRunner */
-    @NonNull public static ProcessRunner getInstance() {
-        if (sProcessRunner == null) {
-            synchronized (ProcessRunner.class) {
-                if (sProcessRunner == null) {
-                    sProcessRunner = new ProcessRunner(
-                            OnDevicePersonalizationApplication.getAppContext(),
-                            new Injector());
-                }
-            }
-        }
-        return sProcessRunner;
-    }
-
+public interface ProcessRunner {
     /** Loads a service in an isolated process */
     @NonNull public ListenableFuture<IsolatedServiceInfo> loadIsolatedService(
-            @NonNull String taskName, @NonNull String packageName) {
-        try {
-            sLogger.d(TAG + ": loadIsolatedService: " + packageName);
-            return loadPlugin(
-                    mInjector.getClock().elapsedRealtime(),
-                    createPluginController(
-                        createPluginId(packageName, taskName),
-                        getPluginManager(mApplicationContext),
-                        packageName));
-        } catch (Exception e) {
-            return Futures.immediateFailedFuture(e);
-        }
-    }
+            @NonNull String taskName, @NonNull String packageName);
 
     /** Executes a service loaded in an isolated process */
     @NonNull public ListenableFuture<Bundle> runIsolatedService(
             @NonNull IsolatedServiceInfo isolatedProcessInfo,
             @NonNull String className,
             int operationCode,
-            @NonNull Bundle serviceParams) {
-        sLogger.d(TAG + ": runIsolatedService: " + className + " op: " + operationCode);
-        Bundle pluginParams = new Bundle();
-        pluginParams.putString(PARAM_CLASS_NAME_KEY, className);
-        pluginParams.putInt(PARAM_OPERATION_KEY, operationCode);
-        pluginParams.putParcelable(PARAM_SERVICE_INPUT, serviceParams);
-        return executePlugin(isolatedProcessInfo.getPluginController(), pluginParams);
-    }
+            @NonNull Bundle serviceParams);
 
     /** Unloads a service loaded in an isolated process */
     @NonNull public ListenableFuture<Void> unloadIsolatedService(
-            @NonNull IsolatedServiceInfo isolatedServiceInfo) {
-        return unloadPlugin(isolatedServiceInfo.getPluginController());
-    }
-
-    @NonNull
-    static PluginManager getPluginManager(@NonNull Context applicationContext) {
-        if (sPluginManager == null) {
-            synchronized (ProcessRunner.class) {
-                if (sPluginManager == null) {
-                    sPluginManager = new PluginManagerImpl(applicationContext);
-                }
-            }
-        }
-        return sPluginManager;
-    }
-
-    @NonNull static PluginController createPluginController(
-            String taskName, @NonNull PluginManager pluginManager, @Nullable String apkName)
-            throws Exception {
-        PluginInfo info = PluginInfo.createJvmInfo(
-                taskName, getArchiveList(apkName), ENTRY_POINT_CLASS);
-        return Objects.requireNonNull(pluginManager.createPluginController(info));
-    }
-
-    @NonNull static ListenableFuture<IsolatedServiceInfo> loadPlugin(
-            long startTimeMillis,
-            @NonNull PluginController pluginController) {
-        return CallbackToFutureAdapter.getFuture(
-            completer -> {
-                try {
-                    sLogger.d(TAG + ": loadPlugin");
-                    pluginController.load(new PluginCallback() {
-                        @Override public void onSuccess(Bundle bundle) {
-                            completer.set(new IsolatedServiceInfo(
-                                    startTimeMillis, pluginController));
-                        }
-                        @Override public void onFailure(FailureType failure) {
-                            completer.setException(new OdpServiceException(
-                                    Constants.STATUS_INTERNAL_ERROR,
-                                    String.format("loadPlugin failed. %s", failure.toString())));
-                        }
-                    });
-                } catch (Exception e) {
-                    completer.setException(e);
-                }
-                return "loadPlugin";
-            }
-        );
-    }
-
-    @NonNull static ListenableFuture<Bundle> executePlugin(
-            @NonNull PluginController pluginController, @NonNull Bundle pluginParams) {
-        return CallbackToFutureAdapter.getFuture(
-            completer -> {
-                try {
-                    sLogger.d(TAG + ": executePlugin");
-                    pluginController.execute(pluginParams, new PluginCallback() {
-                        @Override public void onSuccess(Bundle bundle) {
-                            completer.set(bundle);
-                        }
-                        @Override public void onFailure(FailureType failure) {
-                            completer.setException(new OdpServiceException(
-                                    Constants.STATUS_INTERNAL_ERROR,
-                                    String.format("executePlugin failed: %s", failure.toString())));
-                        }
-                    });
-                } catch (Exception e) {
-                    completer.setException(e);
-                }
-                return "executePlugin";
-            }
-        );
-    }
-
-    @NonNull static ListenableFuture<Void> unloadPlugin(
-            @NonNull PluginController pluginController) {
-        return CallbackToFutureAdapter.getFuture(
-            completer -> {
-                try {
-                    sLogger.d(TAG + ": unloadPlugin");
-                    pluginController.unload(new PluginCallback() {
-                        @Override public void onSuccess(Bundle bundle) {
-                            completer.set(null);
-                        }
-                        @Override public void onFailure(FailureType failure) {
-                            completer.setException(new OdpServiceException(
-                                    Constants.STATUS_INTERNAL_ERROR,
-                                    String.format("executePlugin failed: %s", failure.toString())));
-                        }
-                    });
-                } catch (Exception e) {
-                    completer.setException(e);
-                }
-                return "executePlugin";
-            }
-        );
-    }
-
-    @NonNull static ImmutableList<PluginInfo.ArchiveInfo> getArchiveList(
-            @Nullable String apkName) {
-        if (apkName == null) {
-            return ImmutableList.of();
-        }
-        ImmutableList.Builder<PluginInfo.ArchiveInfo> archiveInfoBuilder = ImmutableList.builder();
-        archiveInfoBuilder.add(
-                PluginInfo.ArchiveInfo.builder().setPackageName(apkName).build());
-        return archiveInfoBuilder.build();
-    }
-
-    static String createPluginId(String vendorPackageName, String taskName) {
-        // TODO(b/249345663) Perform any validation needed on the input.
-        return vendorPackageName + "-" + taskName;
-    }
-}
+            @NonNull IsolatedServiceInfo isolatedServiceInfo);
+}
\ No newline at end of file
diff --git a/src/com/android/ondevicepersonalization/services/process/ProcessRunnerImpl.java b/src/com/android/ondevicepersonalization/services/process/ProcessRunnerImpl.java
new file mode 100644
index 00000000..70422559
--- /dev/null
+++ b/src/com/android/ondevicepersonalization/services/process/ProcessRunnerImpl.java
@@ -0,0 +1,237 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.ondevicepersonalization.services.process;
+
+import android.adservices.ondevicepersonalization.Constants;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.Context;
+import android.os.Bundle;
+
+
+import androidx.concurrent.futures.CallbackToFutureAdapter;
+
+import com.android.ondevicepersonalization.internal.util.LoggerFactory;
+import com.android.ondevicepersonalization.libraries.plugin.FailureType;
+import com.android.ondevicepersonalization.libraries.plugin.PluginCallback;
+import com.android.ondevicepersonalization.libraries.plugin.PluginController;
+import com.android.ondevicepersonalization.libraries.plugin.PluginInfo;
+import com.android.ondevicepersonalization.libraries.plugin.PluginManager;
+import com.android.ondevicepersonalization.libraries.plugin.impl.PluginManagerImpl;
+import com.android.ondevicepersonalization.services.OdpServiceException;
+import com.android.ondevicepersonalization.services.OnDevicePersonalizationApplication;
+import com.android.ondevicepersonalization.services.util.Clock;
+import com.android.ondevicepersonalization.services.util.MonotonicClock;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+
+import java.util.Objects;
+
+/** Utilities to support loading and executing plugins. */
+public class ProcessRunnerImpl implements ProcessRunner {
+    private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
+    private static final String TAG = "ProcessUtils";
+    private static final String ENTRY_POINT_CLASS =
+            "com.android.ondevicepersonalization.services.process.OnDevicePersonalizationPlugin";
+
+    public static final String PARAM_CLASS_NAME_KEY = "param.classname";
+    public static final String PARAM_OPERATION_KEY = "param.operation";
+    public static final String PARAM_SERVICE_INPUT = "param.service_input";
+
+    @NonNull private Context mApplicationContext;
+
+    private static volatile ProcessRunnerImpl sProcessRunner;
+    private static volatile PluginManager sPluginManager;
+
+    static class Injector {
+        Clock getClock() {
+            return MonotonicClock.getInstance();
+        }
+    }
+
+    private final Injector mInjector;
+
+    /** Creates a ProcessRunner. */
+    ProcessRunnerImpl(
+            @NonNull Context applicationContext,
+            @NonNull Injector injector) {
+        mApplicationContext = Objects.requireNonNull(applicationContext);
+        mInjector = Objects.requireNonNull(injector);
+    }
+
+    /** Returns the global ProcessRunner */
+    @NonNull public static ProcessRunnerImpl getInstance() {
+        if (sProcessRunner == null) {
+            synchronized (ProcessRunnerImpl.class) {
+                if (sProcessRunner == null) {
+                    sProcessRunner = new ProcessRunnerImpl(
+                            OnDevicePersonalizationApplication.getAppContext(),
+                            new Injector());
+                }
+            }
+        }
+        return sProcessRunner;
+    }
+
+    /** Loads a service in an isolated process */
+    @Override @NonNull public ListenableFuture<IsolatedServiceInfo> loadIsolatedService(
+            @NonNull String taskName, @NonNull String packageName) {
+        try {
+            sLogger.d(TAG + ": loadIsolatedService: " + packageName);
+            return loadPlugin(
+                    mInjector.getClock().elapsedRealtime(),
+                    createPluginController(
+                        createPluginId(packageName, taskName),
+                        getPluginManager(mApplicationContext),
+                        packageName));
+        } catch (Exception e) {
+            return Futures.immediateFailedFuture(e);
+        }
+    }
+
+    /** Executes a service loaded in an isolated process */
+    @Override @NonNull public ListenableFuture<Bundle> runIsolatedService(
+            @NonNull IsolatedServiceInfo isolatedProcessInfo,
+            @NonNull String className,
+            int operationCode,
+            @NonNull Bundle serviceParams) {
+        sLogger.d(TAG + ": runIsolatedService: " + className + " op: " + operationCode);
+        Bundle pluginParams = new Bundle();
+        pluginParams.putString(PARAM_CLASS_NAME_KEY, className);
+        pluginParams.putInt(PARAM_OPERATION_KEY, operationCode);
+        pluginParams.putParcelable(PARAM_SERVICE_INPUT, serviceParams);
+        return executePlugin(isolatedProcessInfo.getPluginController(), pluginParams);
+    }
+
+    /** Unloads a service loaded in an isolated process */
+    @Override @NonNull public ListenableFuture<Void> unloadIsolatedService(
+            @NonNull IsolatedServiceInfo isolatedServiceInfo) {
+        return unloadPlugin(isolatedServiceInfo.getPluginController());
+    }
+
+    @NonNull
+    static PluginManager getPluginManager(@NonNull Context applicationContext) {
+        if (sPluginManager == null) {
+            synchronized (ProcessRunnerImpl.class) {
+                if (sPluginManager == null) {
+                    sPluginManager = new PluginManagerImpl(applicationContext);
+                }
+            }
+        }
+        return sPluginManager;
+    }
+
+    @NonNull static PluginController createPluginController(
+            String taskName, @NonNull PluginManager pluginManager, @Nullable String apkName)
+            throws Exception {
+        PluginInfo info = PluginInfo.createJvmInfo(
+                taskName, getArchiveList(apkName), ENTRY_POINT_CLASS);
+        return Objects.requireNonNull(pluginManager.createPluginController(info));
+    }
+
+    @NonNull static ListenableFuture<IsolatedServiceInfo> loadPlugin(
+            long startTimeMillis,
+            @NonNull PluginController pluginController) {
+        return CallbackToFutureAdapter.getFuture(
+            completer -> {
+                try {
+                    sLogger.d(TAG + ": loadPlugin");
+                    pluginController.load(new PluginCallback() {
+                        @Override public void onSuccess(Bundle bundle) {
+                            completer.set(new IsolatedServiceInfo(
+                                    startTimeMillis, pluginController));
+                        }
+                        @Override public void onFailure(FailureType failure) {
+                            completer.setException(new OdpServiceException(
+                                    Constants.STATUS_INTERNAL_ERROR,
+                                    String.format("loadPlugin failed. %s", failure.toString())));
+                        }
+                    });
+                } catch (Exception e) {
+                    completer.setException(e);
+                }
+                return "loadPlugin";
+            }
+        );
+    }
+
+    @NonNull static ListenableFuture<Bundle> executePlugin(
+            @NonNull PluginController pluginController, @NonNull Bundle pluginParams) {
+        return CallbackToFutureAdapter.getFuture(
+            completer -> {
+                try {
+                    sLogger.d(TAG + ": executePlugin");
+                    pluginController.execute(pluginParams, new PluginCallback() {
+                        @Override public void onSuccess(Bundle bundle) {
+                            completer.set(bundle);
+                        }
+                        @Override public void onFailure(FailureType failure) {
+                            completer.setException(new OdpServiceException(
+                                    Constants.STATUS_INTERNAL_ERROR,
+                                    String.format("executePlugin failed: %s", failure.toString())));
+                        }
+                    });
+                } catch (Exception e) {
+                    completer.setException(e);
+                }
+                return "executePlugin";
+            }
+        );
+    }
+
+    @NonNull static ListenableFuture<Void> unloadPlugin(
+            @NonNull PluginController pluginController) {
+        return CallbackToFutureAdapter.getFuture(
+            completer -> {
+                try {
+                    sLogger.d(TAG + ": unloadPlugin");
+                    pluginController.unload(new PluginCallback() {
+                        @Override public void onSuccess(Bundle bundle) {
+                            completer.set(null);
+                        }
+                        @Override public void onFailure(FailureType failure) {
+                            completer.setException(new OdpServiceException(
+                                    Constants.STATUS_INTERNAL_ERROR,
+                                    String.format("executePlugin failed: %s", failure.toString())));
+                        }
+                    });
+                } catch (Exception e) {
+                    completer.setException(e);
+                }
+                return "executePlugin";
+            }
+        );
+    }
+
+    @NonNull static ImmutableList<PluginInfo.ArchiveInfo> getArchiveList(
+            @Nullable String apkName) {
+        if (apkName == null) {
+            return ImmutableList.of();
+        }
+        ImmutableList.Builder<PluginInfo.ArchiveInfo> archiveInfoBuilder = ImmutableList.builder();
+        archiveInfoBuilder.add(
+                PluginInfo.ArchiveInfo.builder().setPackageName(apkName).build());
+        return archiveInfoBuilder.build();
+    }
+
+    static String createPluginId(String vendorPackageName, String taskName) {
+        // TODO(b/249345663) Perform any validation needed on the input.
+        return vendorPackageName + "-" + taskName;
+    }
+}
diff --git a/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java b/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
index 88d94214..03ff4b27 100644
--- a/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
+++ b/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
@@ -48,6 +48,7 @@ import com.android.ondevicepersonalization.services.manifest.AppManifestConfigHe
 import com.android.ondevicepersonalization.services.policyengine.UserDataAccessor;
 import com.android.ondevicepersonalization.services.process.IsolatedServiceInfo;
 import com.android.ondevicepersonalization.services.process.ProcessRunner;
+import com.android.ondevicepersonalization.services.process.ProcessRunnerImpl;
 import com.android.ondevicepersonalization.services.statsd.ApiCallStats;
 import com.android.ondevicepersonalization.services.statsd.OdpStatsdLogger;
 import com.android.ondevicepersonalization.services.util.Clock;
@@ -109,7 +110,7 @@ public class AppRequestFlow {
         }
 
         ProcessRunner getProcessRunner() {
-            return ProcessRunner.getInstance();
+            return ProcessRunnerImpl.getInstance();
         }
     }
 
diff --git a/src/com/android/ondevicepersonalization/services/request/RenderFlow.java b/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
index d88e67a5..94facdd8 100644
--- a/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
+++ b/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
@@ -40,6 +40,7 @@ import com.android.ondevicepersonalization.services.display.DisplayHelper;
 import com.android.ondevicepersonalization.services.manifest.AppManifestConfigHelper;
 import com.android.ondevicepersonalization.services.process.IsolatedServiceInfo;
 import com.android.ondevicepersonalization.services.process.ProcessRunner;
+import com.android.ondevicepersonalization.services.process.ProcessRunnerImpl;
 import com.android.ondevicepersonalization.services.statsd.ApiCallStats;
 import com.android.ondevicepersonalization.services.statsd.OdpStatsdLogger;
 import com.android.ondevicepersonalization.services.util.Clock;
@@ -88,7 +89,7 @@ public class RenderFlow {
         }
 
         ProcessRunner getProcessRunner() {
-            return ProcessRunner.getInstance();
+            return ProcessRunnerImpl.getInstance();
         }
     }
 
diff --git a/tests/endtoendtests/OdpClient/res/layout/activity_main.xml b/tests/endtoendtests/OdpClient/res/layout/activity_main.xml
index acee6c76..65cdd23b 100644
--- a/tests/endtoendtests/OdpClient/res/layout/activity_main.xml
+++ b/tests/endtoendtests/OdpClient/res/layout/activity_main.xml
@@ -72,13 +72,6 @@
             style="?android:attr/buttonBarButtonStyle"
             android:text="@string/report_conversion" />
 
-        <Button
-            android:id="@+id/set_status_button"
-            android:layout_height="wrap_content"
-            android:layout_width="wrap_content"
-            style="?android:attr/buttonBarButtonStyle"
-            android:text="@string/set_status" />
-
         <SurfaceView
             android:id="@+id/rendered_view"
             android:layout_width="200dp"
diff --git a/tests/endtoendtests/OdpClient/res/values/strings.xml b/tests/endtoendtests/OdpClient/res/values/strings.xml
index 54799014..a940333e 100644
--- a/tests/endtoendtests/OdpClient/res/values/strings.xml
+++ b/tests/endtoendtests/OdpClient/res/values/strings.xml
@@ -20,5 +20,4 @@
     <string name="get_ad">Get Ad</string>
     <string name="schedule_training">Schedule Training</string>
     <string name="report_conversion">Report Conversion</string>
-    <string name="set_status">Set Status</string>
 </resources>
diff --git a/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java b/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
index c20df2b5..263b61e4 100644
--- a/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
+++ b/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
@@ -16,7 +16,6 @@
 
 package com.example.odpclient;
 
-import android.adservices.ondevicepersonalization.OnDevicePersonalizationConfigManager;
 import android.adservices.ondevicepersonalization.OnDevicePersonalizationManager;
 import android.adservices.ondevicepersonalization.SurfacePackageToken;
 import android.app.Activity;
@@ -37,8 +36,6 @@ import android.widget.Button;
 import android.widget.EditText;
 import android.widget.Toast;
 
-import androidx.annotation.NonNull;
-
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -47,13 +44,11 @@ import java.util.concurrent.atomic.AtomicReference;
 
 public class MainActivity extends Activity {
     private static final String TAG = "OdpClient";
-    private OnDevicePersonalizationConfigManager mOdpConfigManager = null;
 
     private EditText mTextBox;
     private Button mGetAdButton;
     private EditText mScheduleTrainingTextBox;
     private Button mScheduleTrainingButton;
-    private Button mSetStatusButton;
     private EditText mReportConversionTextBox;
     private Button mReportConversionButton;
     private SurfaceView mRenderedView;
@@ -79,23 +74,17 @@ public class MainActivity extends Activity {
         super.onCreate(savedInstanceState);
         setContentView(R.layout.activity_main);
         mContext = getApplicationContext();
-        if (mOdpConfigManager == null) {
-            mOdpConfigManager = mContext.getSystemService(
-                            OnDevicePersonalizationConfigManager.class);
-        }
         mRenderedView = findViewById(R.id.rendered_view);
         mRenderedView.setVisibility(View.INVISIBLE);
         mRenderedView.getHolder().addCallback(new SurfaceCallback());
         mGetAdButton = findViewById(R.id.get_ad_button);
         mScheduleTrainingButton = findViewById(R.id.schedule_training_button);
-        mSetStatusButton = findViewById(R.id.set_status_button);
         mReportConversionButton = findViewById(R.id.report_conversion_button);
         mTextBox = findViewById(R.id.text_box);
         mScheduleTrainingTextBox = findViewById(R.id.schedule_training_text_box);
         mReportConversionTextBox = findViewById(R.id.report_conversion_text_box);
         registerGetAdButton();
         registerScheduleTrainingButton();
-        registerSetStatusButton();
         registerReportConversionButton();
     }
 
@@ -104,10 +93,6 @@ public class MainActivity extends Activity {
                 v -> makeRequest());
     }
 
-    private void registerSetStatusButton() {
-        mSetStatusButton.setOnClickListener(v -> setPersonalizationStatus());
-    }
-
     private void registerReportConversionButton() {
         mReportConversionButton.setOnClickListener(v -> reportConversion());
     }
@@ -258,26 +243,6 @@ public class MainActivity extends Activity {
         }
     }
 
-    private void setPersonalizationStatus() {
-        if (mOdpConfigManager == null) {
-            makeToast("OnDevicePersonalizationConfigManager is null");
-        }
-        boolean enabled = true;
-        mOdpConfigManager.setPersonalizationEnabled(enabled,
-                sCallbackExecutor,
-                new OutcomeReceiver<Void, Exception>() {
-                    @Override
-                    public void onResult(Void result) {
-                        makeToast("Personalization status is set to " + enabled);
-                    }
-
-                    @Override
-                    public void onError(@NonNull Exception error) {
-                        makeToast(error.getMessage());
-                    }
-                });
-    }
-
     private void makeToast(String message) {
         Log.i(TAG, message);
         runOnUiThread(() -> Toast.makeText(MainActivity.this, message, Toast.LENGTH_LONG).show());
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java
index 90340bf9..97651244 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java
@@ -17,6 +17,8 @@
 package com.android.federatedcompute.services.common;
 
 import static com.android.federatedcompute.services.common.Flags.FEDERATED_COMPUTE_GLOBAL_KILL_SWITCH;
+import static com.android.federatedcompute.services.common.Flags.USE_BACKGROUND_ENCRYPTION_KEY_FETCH;
+import static com.android.federatedcompute.services.common.PhFlags.ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH;
 import static com.android.federatedcompute.services.common.PhFlags.FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL;
 import static com.android.federatedcompute.services.common.PhFlags.KEY_FEDERATED_COMPUTE_KILL_SWITCH;
 
@@ -74,4 +76,29 @@ public class PhFlagsTest {
         Flags phFlags = FlagsFactory.getFlags();
         assertThat(phFlags.getEncryptionKeyFetchUrl()).isEqualTo(overrideUrl);
     }
+
+    @Test
+    public void testEnableBackgroundEncryptionKeyFetch() {
+        // Without Overriding
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH,
+                Boolean.toString(USE_BACKGROUND_ENCRYPTION_KEY_FETCH),
+                /* makeDefault= */ false);
+        assertThat(FlagsFactory.getFlags().getEnableBackgroundEncryptionKeyFetch())
+                .isEqualTo(USE_BACKGROUND_ENCRYPTION_KEY_FETCH);
+
+        // Now overriding the value from PH.
+        boolean overrideEnableBackgroundKeyFetch = false;
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH,
+                Boolean.toString(overrideEnableBackgroundKeyFetch),
+                /* makeDefault= */ false);
+
+        Flags phFlags = FlagsFactory.getFlags();
+        assertThat(phFlags.getEnableBackgroundEncryptionKeyFetch())
+                .isEqualTo(overrideEnableBackgroundKeyFetch);
+
+    }
 }
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTestUtil.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTestUtil.java
index b3a7bb64..d4ad1854 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTestUtil.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTestUtil.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.common;
 
+import static com.android.federatedcompute.services.common.Flags.USE_BACKGROUND_ENCRYPTION_KEY_FETCH;
+import static com.android.federatedcompute.services.common.PhFlags.ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH;
 import static com.android.federatedcompute.services.common.PhFlags.KEY_FEDERATED_COMPUTE_KILL_SWITCH;
 
 import android.provider.DeviceConfig;
@@ -59,4 +61,27 @@ public class PhFlagsTestUtil {
                 Boolean.toString(false),
                 /* makeDefault */ false);
     }
+
+    /**
+     * Enable scheduling the background key fetch job.
+     */
+    public static void enableScheduleBackgroundKeyFetchJob() {
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH,
+                Boolean.toString(USE_BACKGROUND_ENCRYPTION_KEY_FETCH),
+                /* makeDefault= */ false);
+    }
+
+    /**
+     * Disable scheduling the background key fetch job.
+     */
+    public static void disableScheduleBackgroundKeyFetchJob() {
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                ENABLE_BACKGROUND_ENCRYPTION_KEY_FETCH,
+                Boolean.toString(false),
+                /* makeDefault= */ false);
+    }
+
 }
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java
index b6f1bbc2..865a8356 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java
@@ -136,6 +136,27 @@ public final class FederatedTrainingTaskDaoTest {
     }
 
     @Test
+    public void findAndRemoveTaskByPopulationNameAndCallingPackage_success() {
+        FederatedTrainingTask task = createDefaultFederatedTrainingTask();
+        mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task);
+        FederatedTrainingTask task2 =
+                createDefaultFederatedTrainingTask().toBuilder()
+                        .jobId(456)
+                        .appPackageName(PACKAGE_NAME)
+                        .populationName(POPULATION_NAME + "_2")
+                        .build();
+        mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task2);
+        assertThat(mTrainingTaskDao.getFederatedTrainingTask(null, null)).hasSize(2);
+
+        FederatedTrainingTask removedTask =
+                mTrainingTaskDao.findAndRemoveTaskByPopulationNameAndCallingPackage(
+                        POPULATION_NAME, PACKAGE_NAME);
+
+        assertThat(DataTestUtil.isEqualTask(removedTask, task)).isTrue();
+        assertThat(mTrainingTaskDao.getFederatedTrainingTask(null, null)).hasSize(1);
+    }
+
+    @Test
     public void findAndRemoveTaskByPopulationName_nonExist() {
         FederatedTrainingTask removedTask =
                 mTrainingTaskDao.findAndRemoveTaskByPopulationName(POPULATION_NAME);
@@ -143,6 +164,15 @@ public final class FederatedTrainingTaskDaoTest {
         assertThat(removedTask).isNull();
     }
 
+    @Test
+    public void findAndRemoveTaskByPopulationNameAndCallingPackage_nonExist() {
+        FederatedTrainingTask removedTask =
+                mTrainingTaskDao.findAndRemoveTaskByPopulationNameAndCallingPackage(
+                        POPULATION_NAME, PACKAGE_NAME);
+
+        assertThat(removedTask).isNull();
+    }
+
     private static byte[] createDefaultTrainingConstraints() {
         FlatBufferBuilder builder = new FlatBufferBuilder();
         builder.finish(TrainingConstraints.createTrainingConstraints(builder, true, true, true));
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java
index 0d1e410c..2a525912 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java
@@ -88,6 +88,7 @@ public class BackgroundKeyFetchJobServiceTest {
     public void setUp() throws Exception {
         PhFlagsTestUtil.setUpDeviceConfigPermissions();
         PhFlagsTestUtil.disableGlobalKillSwitch();
+        PhFlagsTestUtil.enableScheduleBackgroundKeyFetchJob();
         MockitoAnnotations.initMocks(this);
         mContext = ApplicationProvider.getApplicationContext();
         mInjector = new TestInjector();
@@ -190,6 +191,15 @@ public class BackgroundKeyFetchJobServiceTest {
     }
 
     @Test
+    public void testScheduleJob_notEnabled() {
+        PhFlagsTestUtil.disableScheduleBackgroundKeyFetchJob();
+
+        assertThat(BackgroundKeyFetchJobService.scheduleJobIfNeeded(
+                mContext, FlagsFactory.getFlags()
+        )).isEqualTo(false);
+    }
+
+    @Test
     public void testOnStopJob() {
         assertFalse(mSpyService.onStopJob(mock(JobParameters.class)));
     }
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/security/KeyAttestationTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/security/KeyAttestationTest.java
new file mode 100644
index 00000000..bf58be43
--- /dev/null
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/security/KeyAttestationTest.java
@@ -0,0 +1,173 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.federatedcompute.services.security;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+
+import android.security.keystore.KeyProperties;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mockito;
+
+import java.security.InvalidAlgorithmParameterException;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.KeyStore;
+import java.security.KeyStoreException;
+import java.security.cert.CertificateException;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+public final class KeyAttestationTest {
+    private static final byte[] CHALLENGE =
+            ("AHXUDhoSEFikqOefmo8xE7kGp/xjVMRDYBecBiHGxCN8rTv9W0Z4L/14d0OLB"
+                            + "vC1VVzXBAnjgHoKLZzuJifTOaBJwGNIQ2ejnx3n6ayoRchDNCgpK29T+EAhBWzH")
+                    .getBytes();
+
+    private static final String CALLING_APP = "sampleApp1";
+
+    private static final String ANDROID_KEY_STORE = "AndroidKeyStore";
+
+    private static final String KEY_ALIAS = CALLING_APP + "-ODPKeyAttestation";
+
+    private KeyAttestation mKeyAttestation;
+
+    private KeyStore mSpyKeyStore;
+
+    private KeyPairGenerator mSpyKeyPairGenerator;
+
+    @Before
+    public void setUp() throws Exception {
+        mSpyKeyStore = spy(KeyStore.getInstance(ANDROID_KEY_STORE));
+        mSpyKeyPairGenerator = spy(KeyPairGenerator
+                .getInstance(KeyProperties.KEY_ALGORITHM_EC, ANDROID_KEY_STORE));
+        mKeyAttestation = KeyAttestation.getInstanceForTest(
+                ApplicationProvider.getApplicationContext(), new TestInjector());
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        KeyStore keyStore = KeyStore.getInstance(ANDROID_KEY_STORE);
+        keyStore.load(null);
+        if (keyStore.containsAlias(KEY_ALIAS)) {
+            keyStore.deleteEntry(KEY_ALIAS);
+        }
+
+        Mockito.reset(mSpyKeyStore, mSpyKeyPairGenerator);
+    }
+
+    @Test
+    public void testGenerateAttestationRecord_success() {
+        List<String> record = mKeyAttestation.generateAttestationRecord(CHALLENGE, CALLING_APP);
+
+        assertThat(record.size()).isGreaterThan(0);
+    }
+
+    @Test
+    public void testGenerateAttestationRecord_nullKey() {
+        doReturn(null).when(mSpyKeyPairGenerator).generateKeyPair();
+
+        List<String> record = mKeyAttestation.generateAttestationRecord(CHALLENGE, CALLING_APP);
+
+        assertThat(record.size()).isEqualTo(0);
+    }
+
+    @Test
+    public void testGenerateHybridKey_success() {
+        String keyAlias = CALLING_APP + "-ODPKeyAttestation";
+
+        KeyPair keyPair = mKeyAttestation.generateHybridKey(CHALLENGE, keyAlias);
+
+        assertThat(keyPair).isNotNull();
+        assertThat(keyPair.getPublic()).isNotNull();
+        assertThat(keyPair.getPrivate()).isNotNull();
+    }
+
+
+    @Test
+    public void testGenerateHybridKey_initFailure() throws Exception {
+        doThrow(new InvalidAlgorithmParameterException("Invalid Parameters"))
+                .when(mSpyKeyPairGenerator).initialize(any());
+
+        KeyPair keyPair = mKeyAttestation.generateHybridKey(CHALLENGE, KEY_ALIAS);
+
+        assertThat(keyPair).isNull();
+    }
+
+
+    @Test
+    public void testGetAttestationRecordFromKeyAlias_noKey() {
+        String keyAlias2 = CALLING_APP + "-ODPKeyAttestation2";
+
+        KeyPair unused = mKeyAttestation.generateHybridKey(CHALLENGE, KEY_ALIAS);
+        List<String> record = mKeyAttestation.getAttestationRecordFromKeyAlias(keyAlias2);
+
+        assertThat(record.size()).isEqualTo(0);
+    }
+
+    @Test
+    public void testGetAttestationRecordFromKeyAlias_success() {
+        KeyPair unused = mKeyAttestation.generateHybridKey(CHALLENGE, KEY_ALIAS);
+
+        List<String> record = mKeyAttestation.getAttestationRecordFromKeyAlias(KEY_ALIAS);
+
+        assertThat(record.size()).isGreaterThan(0);
+    }
+
+    @Test
+    public void testGetAttestationRecordFromKeyAlias_certFailure() throws Exception {
+        doThrow(new CertificateException("Cert Exception"))
+                .when(mSpyKeyStore).load(any());
+
+        List<String> record = mKeyAttestation.getAttestationRecordFromKeyAlias(KEY_ALIAS);
+
+        assertThat(record.size()).isEqualTo(0);
+    }
+
+    @Test
+    public void testGetAttestationRecordFromKeyAlias_keyStoreFailure() throws Exception {
+        doThrow(new KeyStoreException("Key Store Exception"))
+                .when(mSpyKeyStore).getCertificateChain(any());
+
+        List<String> record = mKeyAttestation.getAttestationRecordFromKeyAlias(KEY_ALIAS);
+
+        assertThat(record.size()).isEqualTo(0);
+    }
+
+    class TestInjector extends KeyAttestation.Injector {
+        @Override
+        KeyStore getKeyStore() {
+            return mSpyKeyStore;
+        }
+
+        @Override
+        KeyPairGenerator getKeyPairGenerator() {
+            return mSpyKeyPairGenerator;
+        }
+    }
+}
diff --git a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleNonExistentPopulationForTraining.java b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleNonExistentPopulationForTraining.java
new file mode 100644
index 00000000..27ee61c7
--- /dev/null
+++ b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleNonExistentPopulationForTraining.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.federatedcompute.test.scenario.federatedcompute;
+
+import android.platform.test.scenario.annotation.Scenario;
+
+import org.junit.AfterClass;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.IOException;
+
+@Scenario
+@RunWith(JUnit4.class)
+/**
+ * Schedule a non-existent population training task from Odp Test app UI
+ * Force the task execution through ADB commands and verify error handling and exit behavior
+ */
+public class ScheduleNonExistentPopulationForTraining {
+    private TestHelper mTestHelper = new TestHelper();
+
+    /** Prepare the device before entering the test class */
+    @BeforeClass
+    public static void prepareDevice() throws IOException {
+        TestHelper.initialize();
+        TestHelper.killRunningProcess();
+    }
+
+    @Before
+    public void setup() throws IOException {
+        mTestHelper.pressHome();
+        mTestHelper.openTestApp();
+        mTestHelper.inputNonExistentPopulationForScheduleTraining();
+    }
+
+    @Test
+    public void testScheduleNonExistentPopulationForTraining() throws IOException {
+        mTestHelper.clickScheduleTraining();
+        mTestHelper.forceExecuteTrainingForNonExistentPopulation();
+    }
+
+    /** Return device to original state after test exeuction */
+    @AfterClass
+    public static void tearDown() throws IOException {
+        TestHelper.pressHome();
+        TestHelper.wrapUp();
+    }
+}
diff --git a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
index 07ec1f87..4c135533 100644
--- a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
+++ b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/TestHelper.java
@@ -34,6 +34,7 @@ public class TestHelper {
     private static UiDevice sUiDevice;
     private static final long UI_FIND_RESOURCE_TIMEOUT = 5000;
     private static final long TRAINING_TASK_COMPLETION_TIMEOUT = 120_000;
+    private static final long CHECKIN_REJECTION_COMPLETION_TIMEOUT = 20_000;
     private static final String ODP_CLIENT_TEST_APP_PACKAGE_NAME = "com.example.odpclient";
     private static final String SCHEDULE_TRAINING_BUTTON_RESOURCE_ID = "schedule_training_button";
     private static final String SCHEDULE_TRAINING_TEXT_BOX_RESOURCE_ID =
@@ -42,6 +43,10 @@ public class TestHelper {
     private static final String ODP_TEST_APP_TRAINING_TASK_JOB_ID = "-630817781";
     private static final String FEDERATED_TRAINING_JOB_SUCCESS_LOG =
             "FederatedJobService - Federated computation job -630817781 is done";
+    private static final String NON_EXISTENT_POPULATION_NAME = "test_non_existent_population";
+    private static final String NON_EXISTENT_POPULATION_TASK_JOB_ID = "1892833995";
+    private static final String NON_EXISTENT_POPULATION_JOB_FAILURE_LOG =
+            "job 1892833995 was rejected during check in, reason NO_TASK_AVAILABLE";
 
     public static void pressHome() {
         getUiDevice().pressHome();
@@ -119,6 +124,13 @@ public class TestHelper {
         scheduleTrainingTextBox.setText(ODP_TEST_APP_POPULATION_NAME);
     }
 
+    /** Put a test non existent population name down for training */
+    public void inputNonExistentPopulationForScheduleTraining() {
+        UiObject2 scheduleTrainingTextBox = getScheduleTrainingTextBox();
+        assertNotNull("Schedule Training text box not found", scheduleTrainingTextBox);
+        scheduleTrainingTextBox.setText(NON_EXISTENT_POPULATION_NAME);
+    }
+
     /** Click Schedule Training button. */
     public void clickScheduleTraining() {
         UiObject2 scheduleTrainingButton = getScheduleTrainingButton();
@@ -143,11 +155,34 @@ public class TestHelper {
 
         if (!foundTrainingJobSuccessLog) {
             Assert.fail(String.format(
-                    "Failed to find federated training job success log within test window %d ms",
+                    "Failed to find federated training job success log %s within test window %d ms",
+                    FEDERATED_TRAINING_JOB_SUCCESS_LOG,
                     TRAINING_TASK_COMPLETION_TIMEOUT));
         }
     }
 
+    /** Force the JobScheduler to execute the training task for non existent population */
+    public void forceExecuteTrainingForNonExistentPopulation() throws IOException {
+        executeShellCommand("logcat -c"); // Cleans the log buffer
+        executeShellCommand("logcat -G 32M"); // Set log buffer to 32MB
+        executeShellCommand(
+                "cmd jobscheduler run -f com.google.android.federatedcompute "
+                        + NON_EXISTENT_POPULATION_TASK_JOB_ID);
+        SystemClock.sleep(10000);
+
+        boolean foundTrainingFailureLog = findLog(
+                NON_EXISTENT_POPULATION_JOB_FAILURE_LOG,
+                CHECKIN_REJECTION_COMPLETION_TIMEOUT,
+                5000);
+
+        if (!foundTrainingFailureLog) {
+            Assert.fail(String.format(
+                    "Failed to find federated training failure log: %s within test window %d ms",
+                    NON_EXISTENT_POPULATION_JOB_FAILURE_LOG,
+                    CHECKIN_REJECTION_COMPLETION_TIMEOUT));
+        }
+    }
+
     /** Attempt to find a specific log entry within the timeout window */
     private boolean findLog(final String targetLog, long timeoutMillis,
             long queryIntervalMillis) throws IOException {
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadHelper.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadHelper.java
index d181f82f..3437c6a1 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadHelper.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadHelper.java
@@ -27,23 +27,19 @@ import java.io.IOException;
 /** Helper class for interacting with download flow. */
 public class DownloadHelper {
     private static UiDevice sUiDevice;
-
-    // Download Url in APK points to a local location
-    private static final String VENDOR_APK_LOCAL_URL_PATH =
-            "data/temp_files/OdpSampleNetworkLocal.apk";
-    // Download Url in APK points to a public url with file size ~ 100KB
-    private static final String VENDOR_APK_SERVER_SMALL_PATH =
-            "data/temp_files/OdpSampleNetworkSmall.apk";
-    // Download Url in APK points to a public url with file size ~ 1MB
-    private static final String VENDOR_APK_SERVER_MEDIUM_PATH =
-            "data/temp_files/OdpSampleNetworkMedium.apk";
-    // Download Url in APK points to a public url with file size ~ 10MB
-    private static final String VENDOR_APK_SERVER_LARGE_PATH =
-            "data/temp_files/OdpSampleNetworkLarge.apk";
-
     private static final String MDD_WIFI_CHARGING_PERIODIC_TASK_JOB_ID = "1003";
     private static final String DOWNLOAD_PROCESSING_TASK_JOB_ID = "1004";
     private static final String MAINTENANCE_TASK_JOB_ID = "1005";
+    private static final String DOWNLOAD_JOB_SUCCESS_LOG =
+            "MddJobService: MddJobService.MddHandleTask succeeded!";
+    private static final String PLUGIN_CODE_COMPLETED_LOG =
+            "OnDevicePersonalizationDataProcessingAsyncCallable: "
+                    + "Plugin filter code completed successfully";
+    private static final String FILTER_AND_STORE_DATA_SUCCESS_LOG =
+            "OnDevicePersonalizationDataProcessingAsyncCallable: "
+                    + "filter and store data completed, transaction successful: true";
+    private static final long DOWNLOAD_JOB_COMPLETION_TIMEOUT = 120_000;
+    private static final long DOWNLOAD_PROCESSING_JOB_COMPLETION_TIMEOUT = 120_000;
 
     /** Commands to prepare the device and odp module before testing. */
     public static void initialize() throws IOException {
@@ -86,30 +82,6 @@ public class DownloadHelper {
         getUiDevice().pressHome();
     }
 
-    public void installVendorApkWithLocalUrl() throws IOException {
-        installVendorApk(VENDOR_APK_LOCAL_URL_PATH);
-    }
-
-    public void installVendorApkWithServerSmallFile() throws IOException {
-        installVendorApk(VENDOR_APK_SERVER_SMALL_PATH);
-    }
-
-    public void installVendorApkWithServerMediumFile() throws IOException {
-        installVendorApk(VENDOR_APK_SERVER_MEDIUM_PATH);
-    }
-
-    public void installVendorApkWithServerLargeFile() throws IOException {
-        installVendorApk(VENDOR_APK_SERVER_LARGE_PATH);
-    }
-
-    private void installVendorApk(String path) throws IOException {
-        executeShellCommand("pm install -r " + path);
-    }
-
-    public void uninstallVendorApk() throws IOException {
-        executeShellCommand("pm uninstall com.android.odpsamplenetwork");
-    }
-
     public void cleanupDatabase() throws IOException {
         executeShellCommand(
                 "cmd jobscheduler run -f com.google.android.ondevicepersonalization.services "
@@ -125,17 +97,67 @@ public class DownloadHelper {
     }
 
     public void downloadVendorData() throws IOException {
+        executeShellCommand("logcat -c"); // Cleans the log buffer
+        executeShellCommand("logcat -G 32M"); // Set log buffer to 32MB
         executeShellCommand(
                 "cmd jobscheduler run -f com.google.android.ondevicepersonalization.services "
                         + MDD_WIFI_CHARGING_PERIODIC_TASK_JOB_ID);
-        SystemClock.sleep(60000);
+
+        boolean foundDownloadJobSuccessLog = findLog(
+                DOWNLOAD_JOB_SUCCESS_LOG,
+                DOWNLOAD_JOB_COMPLETION_TIMEOUT,
+                5000);
+
+        if (!foundDownloadJobSuccessLog) {
+            Assert.fail(String.format(
+                    "Failed to find download job success log %s within test window %d ms",
+                    DOWNLOAD_JOB_SUCCESS_LOG,
+                    DOWNLOAD_JOB_COMPLETION_TIMEOUT));
+        }
     }
 
     public void processDownloadedVendorData() throws IOException {
         executeShellCommand(
                 "cmd jobscheduler run -f com.google.android.ondevicepersonalization.services "
                         + DOWNLOAD_PROCESSING_TASK_JOB_ID);
-        SystemClock.sleep(5000);
+
+        boolean foundPlugInCodeCompletionLog = findLog(
+                PLUGIN_CODE_COMPLETED_LOG,
+                DOWNLOAD_PROCESSING_JOB_COMPLETION_TIMEOUT,
+                5000);
+
+        if (!foundPlugInCodeCompletionLog) {
+            Assert.fail(String.format(
+                    "Failed to find plugin code completion log %s within test window %d ms",
+                    PLUGIN_CODE_COMPLETED_LOG,
+                    DOWNLOAD_PROCESSING_JOB_COMPLETION_TIMEOUT));
+        }
+
+        boolean foundFilterAndStoreDataSuccessLog = findLog(
+                FILTER_AND_STORE_DATA_SUCCESS_LOG,
+                DOWNLOAD_PROCESSING_JOB_COMPLETION_TIMEOUT,
+                5000);
+
+        if (!foundFilterAndStoreDataSuccessLog) {
+            Assert.fail(String.format(
+                    "Failed to find filter and store data success log %s within test window %d ms",
+                    foundFilterAndStoreDataSuccessLog,
+                    DOWNLOAD_PROCESSING_JOB_COMPLETION_TIMEOUT));
+        }
+    }
+
+    /** Attempt to find a specific log entry within the timeout window */
+    private boolean findLog(final String targetLog, long timeoutMillis,
+                            long queryIntervalMillis) throws IOException {
+
+        long startTime = System.currentTimeMillis();
+        while (System.currentTimeMillis() - startTime < timeoutMillis) {
+            if (getUiDevice().executeShellCommand("logcat -d").contains(targetLog)) {
+                return true;
+            }
+            SystemClock.sleep(queryIntervalMillis);
+        }
+        return false;
     }
 
     private static void executeShellCommand(String cmd) {
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorData.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorData.java
index 2a9f3526..a7585207 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorData.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorData.java
@@ -51,7 +51,6 @@ public class DownloadVendorData {
 
     @After
     public void tearDown() throws IOException {
-        mDownloadHelper.uninstallVendorApk();
         mDownloadHelper.cleanupDatabase();
         mDownloadHelper.cleanupDownloadedMetadata();
         mDownloadHelper.pressHome();
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrl.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrl.java
deleted file mode 100644
index 57d9f9a7..00000000
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrl.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.scenario.annotation.Scenario;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-import java.io.IOException;
-
-@Scenario
-@RunWith(JUnit4.class)
-public class DownloadVendorDataFromLocalUrl {
-
-    private DownloadHelper mDownloadHelper = new DownloadHelper();
-
-    @Before
-    public void setUp() throws IOException {
-        mDownloadHelper.pressHome();
-        mDownloadHelper.installVendorApkWithLocalUrl();
-    }
-
-    @Test
-    public void testDownloadVendorDataFromLocalUrl() throws IOException {
-        mDownloadHelper.downloadVendorData();
-        mDownloadHelper.processDownloadedVendorData();
-    }
-
-    @After
-    public void tearDown() throws IOException {
-        mDownloadHelper.uninstallVendorApk();
-        mDownloadHelper.cleanupDatabase();
-        mDownloadHelper.cleanupDownloadedMetadata();
-        mDownloadHelper.pressHome();
-    }
-
-}
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFile.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFile.java
deleted file mode 100644
index 5aa6852f..00000000
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFile.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.scenario.annotation.Scenario;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-import java.io.IOException;
-
-@Scenario
-@RunWith(JUnit4.class)
-public class DownloadVendorDataFromServerLargeFile {
-
-    private DownloadHelper mDownloadHelper = new DownloadHelper();
-
-    @Before
-    public void setUp() throws IOException {
-        mDownloadHelper.pressHome();
-        mDownloadHelper.installVendorApkWithServerLargeFile();
-    }
-
-    @Test
-    public void testDownloadVendorDataFromServerLargeFile() throws IOException {
-        mDownloadHelper.downloadVendorData();
-        mDownloadHelper.processDownloadedVendorData();
-    }
-
-    @After
-    public void tearDown() throws IOException {
-        mDownloadHelper.uninstallVendorApk();
-        mDownloadHelper.cleanupDatabase();
-        mDownloadHelper.cleanupDownloadedMetadata();
-        mDownloadHelper.pressHome();
-    }
-
-}
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFile.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFile.java
deleted file mode 100644
index fbc67093..00000000
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFile.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.scenario.annotation.Scenario;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-import java.io.IOException;
-
-@Scenario
-@RunWith(JUnit4.class)
-public class DownloadVendorDataFromServerMediumFile {
-
-    private DownloadHelper mDownloadHelper = new DownloadHelper();
-
-    @Before
-    public void setUp() throws IOException {
-        mDownloadHelper.pressHome();
-        mDownloadHelper.installVendorApkWithServerMediumFile();
-    }
-
-    @Test
-    public void testDownloadVendorDataFromServerMediumFile() throws IOException {
-        mDownloadHelper.downloadVendorData();
-        mDownloadHelper.processDownloadedVendorData();
-    }
-
-    @After
-    public void tearDown() throws IOException {
-        mDownloadHelper.uninstallVendorApk();
-        mDownloadHelper.cleanupDatabase();
-        mDownloadHelper.cleanupDownloadedMetadata();
-        mDownloadHelper.pressHome();
-    }
-
-}
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerSmallFile.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerSmallFile.java
deleted file mode 100644
index 9c8a4e2b..00000000
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerSmallFile.java
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.scenario.annotation.Scenario;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-
-import java.io.IOException;
-
-@Scenario
-@RunWith(JUnit4.class)
-public class DownloadVendorDataFromServerSmallFile {
-
-    private DownloadHelper mDownloadHelper = new DownloadHelper();
-
-    @Before
-    public void setUp() throws IOException {
-        mDownloadHelper.pressHome();
-        mDownloadHelper.installVendorApkWithServerSmallFile();
-    }
-
-    @Test
-    public void testDownloadVendorDataFromServerSmallFile() throws IOException {
-        mDownloadHelper.downloadVendorData();
-        mDownloadHelper.processDownloadedVendorData();
-    }
-
-    @After
-    public void tearDown() throws IOException {
-        mDownloadHelper.uninstallVendorApk();
-        mDownloadHelper.cleanupDatabase();
-        mDownloadHelper.cleanupDownloadedMetadata();
-        mDownloadHelper.pressHome();
-    }
-
-}
diff --git a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerSmallFileMicrobenchmark.java b/tests/perftests/scenarios/tests/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleNonExistentPopulationForTrainingMicrobenchmark.java
index 482fc494..8150bf3d 100644
--- a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerSmallFileMicrobenchmark.java
+++ b/tests/perftests/scenarios/tests/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleNonExistentPopulationForTrainingMicrobenchmark.java
@@ -14,10 +14,11 @@
  * limitations under the License.
  */
 
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
+package android.federatedcompute.test.scenario.federatedcompute;
 
 import android.platform.test.microbenchmark.Microbenchmark;
 import android.platform.test.rule.DropCachesRule;
+import android.platform.test.rule.KillAppsRule;
 import android.platform.test.rule.PressHomeRule;
 
 import org.junit.Rule;
@@ -25,10 +26,11 @@ import org.junit.rules.RuleChain;
 import org.junit.runner.RunWith;
 
 @RunWith(Microbenchmark.class)
-public class DownloadVendorDataFromServerSmallFileMicrobenchmark
-        extends DownloadVendorDataFromServerSmallFile {
+public class ScheduleNonExistentPopulationForTrainingMicrobenchmark
+        extends ScheduleNonExistentPopulationForTraining {
 
     @Rule
     public RuleChain rules = RuleChain.outerRule(new DropCachesRule())
+            .around(new KillAppsRule("com.example.odpclient"))
             .around(new PressHomeRule());
 }
diff --git a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrlMicrobenchmark.java b/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrlMicrobenchmark.java
deleted file mode 100644
index 12702042..00000000
--- a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromLocalUrlMicrobenchmark.java
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.microbenchmark.Microbenchmark;
-import android.platform.test.rule.DropCachesRule;
-import android.platform.test.rule.PressHomeRule;
-
-import org.junit.Rule;
-import org.junit.rules.RuleChain;
-import org.junit.runner.RunWith;
-
-@RunWith(Microbenchmark.class)
-public class DownloadVendorDataFromLocalUrlMicrobenchmark extends DownloadVendorDataFromLocalUrl {
-
-    @Rule
-    public RuleChain rules = RuleChain.outerRule(new DropCachesRule())
-            .around(new PressHomeRule());
-}
diff --git a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFileMicrobenchmark.java b/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFileMicrobenchmark.java
deleted file mode 100644
index 83345108..00000000
--- a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerLargeFileMicrobenchmark.java
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.microbenchmark.Microbenchmark;
-import android.platform.test.rule.DropCachesRule;
-import android.platform.test.rule.PressHomeRule;
-
-import org.junit.Rule;
-import org.junit.rules.RuleChain;
-import org.junit.runner.RunWith;
-
-@RunWith(Microbenchmark.class)
-public class DownloadVendorDataFromServerLargeFileMicrobenchmark
-        extends DownloadVendorDataFromServerLargeFile {
-
-    @Rule
-    public RuleChain rules = RuleChain.outerRule(new DropCachesRule())
-            .around(new PressHomeRule());
-}
diff --git a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFileMicrobenchmark.java b/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFileMicrobenchmark.java
deleted file mode 100644
index aa9866ce..00000000
--- a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/DownloadVendorDataFromServerMediumFileMicrobenchmark.java
+++ /dev/null
@@ -1,34 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package android.ondevicepersonalization.test.scenario.ondevicepersonalization;
-
-import android.platform.test.microbenchmark.Microbenchmark;
-import android.platform.test.rule.DropCachesRule;
-import android.platform.test.rule.PressHomeRule;
-
-import org.junit.Rule;
-import org.junit.rules.RuleChain;
-import org.junit.runner.RunWith;
-
-@RunWith(Microbenchmark.class)
-public class DownloadVendorDataFromServerMediumFileMicrobenchmark
-        extends DownloadVendorDataFromServerMediumFile {
-
-    @Rule
-    public RuleChain rules = RuleChain.outerRule(new DropCachesRule())
-            .around(new PressHomeRule());
-}
diff --git a/tests/servicetests/src/com/android/ondevicepersonalization/services/process/ProcessRunnerTest.java b/tests/servicetests/src/com/android/ondevicepersonalization/services/process/ProcessRunnerImplTest.java
index 7f7bdb66..80679c41 100644
--- a/tests/servicetests/src/com/android/ondevicepersonalization/services/process/ProcessRunnerTest.java
+++ b/tests/servicetests/src/com/android/ondevicepersonalization/services/process/ProcessRunnerImplTest.java
@@ -30,18 +30,18 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 @RunWith(JUnit4.class)
-public class ProcessRunnerTest {
-    ProcessRunner mProcessRunner = new ProcessRunner(
+public class ProcessRunnerImplTest {
+    ProcessRunner mProcessRunner = new ProcessRunnerImpl(
             ApplicationProvider.getApplicationContext(),
-            new ProcessRunner.Injector());
+            new ProcessRunnerImpl.Injector());
     @Test
     public void testGetArchiveList_NullApkList() throws Exception {
-        assertTrue(ProcessRunner.getArchiveList(null).isEmpty());
+        assertTrue(ProcessRunnerImpl.getArchiveList(null).isEmpty());
     }
 
     @Test
     public void testGetArchiveList() throws Exception {
-        ImmutableList<PluginInfo.ArchiveInfo> result = ProcessRunner.getArchiveList("fakeApk");
+        ImmutableList<PluginInfo.ArchiveInfo> result = ProcessRunnerImpl.getArchiveList("fakeApk");
         assertEquals(1, result.size());
         assertEquals("fakeApk", result.get(0).packageName());
     }