Snap for 11174750 from 15e3ef7b1a426964cf2c2e080a8cbbd02112f6d0 to mainline-adservices-releaseaml_ads_341413000android14-mainline-adservices-release

Change-Id: I078c52b0d8be683710981e93e17442a0d0a91df8

97 files changed, 5017 insertions, 1787 deletions

diff --git a/federatedcompute/apk/AndroidManifest.xml b/federatedcompute/apk/AndroidManifest.xml
index dec46ab3..652dc05a 100644
--- a/federatedcompute/apk/AndroidManifest.xml
+++ b/federatedcompute/apk/AndroidManifest.xml
@@ -48,6 +48,10 @@
                  android:exported="false"
                  android:permission="android.permission.BIND_JOB_SERVICE">
         </service>
+        <service android:name="com.android.federatedcompute.services.encryption.BackgroundKeyFetchJobService"
+            android:exported="false"
+            android:permission="android.permission.BIND_JOB_SERVICE">
+        </service>
         <service android:name="com.android.federatedcompute.services.training.IsolatedTrainingService"
                  android:isolatedProcess="true" android:exported="false" >
         </service>
diff --git a/federatedcompute/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImpl.java b/federatedcompute/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImpl.java
index 7d8db4d8..a72819b7 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImpl.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImpl.java
@@ -20,6 +20,9 @@ import android.app.Service;
 import android.content.Intent;
 import android.os.IBinder;
 
+import com.android.federatedcompute.services.common.Flags;
+import com.android.federatedcompute.services.common.FlagsFactory;
+import com.android.federatedcompute.services.encryption.BackgroundKeyFetchJobService;
 import com.android.federatedcompute.services.statsd.FederatedComputeStatsdLogger;
 
 import java.util.Objects;
@@ -28,6 +31,12 @@ import java.util.Objects;
 public class FederatedComputeManagingServiceImpl extends Service {
     private FederatedComputeManagingServiceDelegate mFcpServiceDelegate;
 
+    private Flags mFlags;
+
+    public FederatedComputeManagingServiceImpl() {
+        mFlags = FlagsFactory.getFlags();
+    }
+
     @Override
     public void onCreate() {
         super.onCreate();
@@ -35,6 +44,7 @@ public class FederatedComputeManagingServiceImpl extends Service {
             mFcpServiceDelegate =
                     new FederatedComputeManagingServiceDelegate(
                             this, FederatedComputeStatsdLogger.getInstance());
+            BackgroundKeyFetchJobService.scheduleJobIfNeeded(this, mFlags);
         }
     }
 
diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java b/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
index 33b64603..6c697525 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/Constants.java
@@ -36,5 +36,14 @@ public class Constants {
     public static final String ISOLATED_TRAINING_SERVICE_NAME =
             "com.android.federatedcompute.services.training.IsolatedTrainingService";
 
+    public static final String TRACE_HTTP_ISSUE_CHECKIN = "Http#issueCheckin";
+    public static final String TRACE_HTTP_REPORT_RESULT = "Http#reportResult";
+    public static final String TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING =
+            "IsolatedProcess#runFlTraining";
+    public static final String TRACE_NATIVE_RUN_FEDERATED_COMPUTATION =
+            "Native#runFederatedComputation";
+    public static final String TRACE_WORKER_RUN_FL_COMPUTATION = "Worker#runFlComputation";
+    public static final String TRACE_WORKER_START_TRAINING_RUN = "Worker#startTrainingRun";
+
     private Constants() {}
 }
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.aidl b/federatedcompute/src/com/android/federatedcompute/services/common/FederatedComputeJobInfo.java
index 7d841a50..7df5d821 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.aidl
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/FederatedComputeJobInfo.java
@@ -14,6 +14,12 @@
  * limitations under the License.
  */
 
-package android.adservices.ondevicepersonalization;
+package com.android.federatedcompute.services.common;
 
-parcelable TrainingExampleInput;
+public class FederatedComputeJobInfo {
+
+    private FederatedComputeJobInfo() {}
+
+    /** JOB ID to periodically download encryption key. */
+    public static final int ENCRYPTION_KEY_FETCH_JOB_ID = 1000;
+}
diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java b/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
index 5c01501e..6af7c1ca 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/Flags.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.common;
 
+import java.util.concurrent.TimeUnit;
+
 /** FederatedCompute feature flags interface. This Flags interface hold the default values */
 public interface Flags {
     /**
@@ -29,7 +31,10 @@ public interface Flags {
         return FEDERATED_COMPUTE_GLOBAL_KILL_SWITCH;
     }
 
-    /** Flags for {@link FederatedComputeJobManager}. */
+    /**
+     * Flags for {@link
+     * com.android.federatedcompute.services.scheduling.FederatedComputeJobManager}.
+     */
     long DEFAULT_SCHEDULING_PERIOD_SECS = 60 * 5; // 5 minutes
 
     default long getDefaultSchedulingPeriodSecs() {
@@ -80,7 +85,7 @@ public interface Flags {
         return TRANSIENT_ERROR_RETRY_DELAY_SECS;
     }
 
-    /** Flags for {@link FederatedExampleIterator}. */
+    /** Flags for ExampleStoreService. */
     long APP_HOSTED_EXAMPLE_STORE_TIMEOUT_SECS = 30;
 
     default long getAppHostedExampleStoreTimeoutSecs() {
@@ -135,4 +140,30 @@ public interface Flags {
     default long getTrainingConditionCheckThrottlePeriodMillis() {
         return TRAINING_CONDITION_CHECK_THROTTLE_PERIOD_MILLIS;
     }
+
+    String ENCRYPTION_KEY_FETCH_URL =
+            "https://fake-coordinator/v1alpha/publicKeys";
+
+    /**
+     * @return Url to fetch encryption key for federated compute.
+     */
+    default String getEncryptionKeyFetchUrl() {
+        return ENCRYPTION_KEY_FETCH_URL;
+    }
+
+    Long FEDERATED_COMPUTE_ENCRYPTION_KEY_MAX_AGE_SECONDS =
+            TimeUnit.DAYS.toSeconds(14/* duration= */);
+
+    /**
+     * @return default max age in seconds for federated compute ecryption keys.
+     */
+    default Long getFederatedComputeEncryptionKeyMaxAgeSeconds() {
+        return FEDERATED_COMPUTE_ENCRYPTION_KEY_MAX_AGE_SECONDS;
+    }
+
+    Long ENCRYPTION_KEY_FETCH_PERIOD_SECONDS = 60 * 60 * 24L; // every 24 h
+
+    default Long getEncryptionKeyFetchPeriodSeconds() {
+        return ENCRYPTION_KEY_FETCH_PERIOD_SECONDS;
+    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java b/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
index 1072965b..ade34909 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/common/PhFlags.java
@@ -35,6 +35,9 @@ public final class PhFlags implements Flags {
 
     // OnDevicePersonalization Namespace String from DeviceConfig class
     static final String NAMESPACE_ON_DEVICE_PERSONALIZATION = "on_device_personalization";
+
+    static final String FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL =
+            "fcp_encryption_key_download_url";
     private static final PhFlags sSingleton = new PhFlags();
 
     /** Returns the singleton instance of the PhFlags. */
@@ -60,4 +63,13 @@ public final class PhFlags implements Flags {
     static String getSystemPropertyName(String key) {
         return SYSTEM_PROPERTY_PREFIX + key;
     }
+
+    @Override
+    public String getEncryptionKeyFetchUrl() {
+        return DeviceConfig.getString(
+                /* namespace= */ NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                /* name= */ FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL,
+                /* defaultValue= */ ENCRYPTION_KEY_FETCH_URL
+        );
+    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedComputeDbHelper.java b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedComputeDbHelper.java
index 92ff384c..1ed49974 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedComputeDbHelper.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedComputeDbHelper.java
@@ -54,7 +54,7 @@ public class FederatedComputeDbHelper extends SQLiteOpenHelper {
                     + FederatedTrainingTaskColumns.CONTEXT_DATA
                     + " BLOB, "
                     + FederatedTrainingTaskColumns.CREATION_TIME
-                    + " INTEGER, "
+                    + " INTEGER NOT NULL, "
                     + FederatedTrainingTaskColumns.LAST_SCHEDULED_TIME
                     + " INTEGER, "
                     + FederatedTrainingTaskColumns.LAST_RUN_START_TIME
@@ -62,11 +62,14 @@ public class FederatedComputeDbHelper extends SQLiteOpenHelper {
                     + FederatedTrainingTaskColumns.LAST_RUN_END_TIME
                     + " INTEGER, "
                     + FederatedTrainingTaskColumns.EARLIEST_NEXT_RUN_TIME
-                    + " INTEGER, "
+                    + " INTEGER NOT NULL, "
                     + FederatedTrainingTaskColumns.CONSTRAINTS
                     + " BLOB, "
                     + FederatedTrainingTaskColumns.SCHEDULING_REASON
-                    + " INTEGER )";
+                    + " INTEGER, "
+                    + "UNIQUE("
+                    + FederatedTrainingTaskColumns.JOB_SCHEDULER_JOB_ID
+                    + "))";
 
     private static final String CREATE_ENCRYPTION_KEY_TABLE =
             "CREATE TABLE "
diff --git a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
index 409505f7..69ca70a0 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDao.java
@@ -21,6 +21,7 @@ import static com.android.federatedcompute.services.data.FederatedTraningTaskCon
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
+import android.database.SQLException;
 import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteException;
 import android.database.sqlite.SQLiteOpenHelper;
@@ -84,11 +85,20 @@ public class FederatedTrainingTaskDao {
 
     /** Insert a training task or update it if task already exists. */
     public boolean updateOrInsertFederatedTrainingTask(FederatedTrainingTask trainingTask) {
-        SQLiteDatabase db = getWritableDatabase();
-        if (db == null) {
-            throw new SQLiteException("Failed to open database.");
+        try {
+            SQLiteDatabase db = getWritableDatabase();
+            if (db == null) {
+                return false;
+            }
+            return trainingTask.addToDatabase(db);
+        } catch (SQLException e) {
+            LogUtil.e(
+                    TAG,
+                    e,
+                    "Failed to persist federated training task %s",
+                    trainingTask.populationName());
+            return false;
         }
-        return trainingTask.addToDatabase(db);
     }
 
     /** Get the list of tasks that match select conditions. */
@@ -109,10 +119,15 @@ public class FederatedTrainingTaskDao {
         String[] selectionArgs = selectionArgs(jobId);
         FederatedTrainingTask task =
                 Iterables.getOnlyElement(getFederatedTrainingTask(selection, selectionArgs), null);
-        if (task != null) {
-            deleteFederatedTrainingTask(selection, selectionArgs);
+        try {
+            if (task != null) {
+                deleteFederatedTrainingTask(selection, selectionArgs);
+            }
+            return task;
+        } catch (SQLException e) {
+            LogUtil.e(TAG, e, "Failed to delete federated training task by job id %d", jobId);
+            return null;
         }
-        return task;
     }
 
     /** Delete a task from table based on population name. */
@@ -121,10 +136,19 @@ public class FederatedTrainingTaskDao {
         String[] selectionArgs = {populationName};
         FederatedTrainingTask task =
                 Iterables.getOnlyElement(getFederatedTrainingTask(selection, selectionArgs), null);
-        if (task != null) {
-            deleteFederatedTrainingTask(selection, selectionArgs);
+        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",
+                    populationName);
+            return null;
         }
-        return task;
     }
 
     /** Delete a task from table based on population name and job scheduler id. */
@@ -138,10 +162,20 @@ public class FederatedTrainingTaskDao {
         String[] selectionArgs = {populationName, String.valueOf(jobId)};
         FederatedTrainingTask task =
                 Iterables.getOnlyElement(getFederatedTrainingTask(selection, selectionArgs), null);
-        if (task != null) {
-            deleteFederatedTrainingTask(selection, selectionArgs);
+        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 job id %d",
+                    populationName,
+                    jobId);
+            return null;
         }
-        return task;
     }
 
     private String[] selectionArgs(Number... args) {
diff --git a/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java b/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java
new file mode 100644
index 00000000..6b8287ae
--- /dev/null
+++ b/federatedcompute/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobService.java
@@ -0,0 +1,174 @@
+/*
+ * 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.encryption;
+
+import android.app.job.JobInfo;
+import android.app.job.JobParameters;
+import android.app.job.JobScheduler;
+import android.app.job.JobService;
+import android.content.ComponentName;
+import android.content.Context;
+
+import com.android.federatedcompute.internal.util.LogUtil;
+import com.android.federatedcompute.services.common.FederatedComputeExecutors;
+import com.android.federatedcompute.services.common.FederatedComputeJobInfo;
+import com.android.federatedcompute.services.common.Flags;
+import com.android.federatedcompute.services.common.FlagsFactory;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKey;
+import com.android.internal.annotations.VisibleForTesting;
+
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListeningExecutorService;
+
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+
+public class BackgroundKeyFetchJobService extends JobService {
+    private static final String TAG = BackgroundKeyFetchJobService.class.getSimpleName();
+
+    private static final int ENCRYPTION_KEY_FETCH_JOB_ID =
+            FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID;
+
+    static class Injector {
+        ListeningExecutorService getExecutor() {
+            return FederatedComputeExecutors.getBackgroundExecutor();
+        }
+
+        ListeningExecutorService getLightWeightExecutor() {
+            return FederatedComputeExecutors.getLightweightExecutor();
+        }
+
+        FederatedComputeEncryptionKeyManager getEncryptionKeyManager(Context context) {
+            return FederatedComputeEncryptionKeyManager.getInstance(context);
+        }
+    }
+
+    private final Injector mInjector;
+
+    @VisibleForTesting
+    BackgroundKeyFetchJobService(Injector injector) {
+        mInjector = injector;
+    }
+
+    /** Runs the background key fetch and persist keys job. The method is for testing only. */
+    @VisibleForTesting
+    public void run(JobParameters params) {
+        var unused = Futures.submit(() -> onStartJob(params), mInjector.getExecutor());
+    }
+
+    @Override
+    public boolean onStartJob(JobParameters params) {
+        LogUtil.d(TAG, "BackgroundKeyFetchJobService.onStartJob %d", params.getJobId());
+        if (FlagsFactory.getFlags().getGlobalKillSwitch()) {
+            LogUtil.d(TAG, "GlobalKillSwitch enabled, finishing job.");
+            jobFinished(params, false /* wantsReschedule= */);
+            return true;
+        }
+        mInjector
+                .getEncryptionKeyManager(this)
+                .fetchAndPersistActiveKeys(FederatedComputeEncryptionKey.KEY_TYPE_ENCRYPTION,
+                        /* isScheduledJob= */ true)
+                .addCallback(
+                        new FutureCallback<List<FederatedComputeEncryptionKey>>() {
+                            @Override
+                            public void onSuccess(
+                                    List<FederatedComputeEncryptionKey>
+                                            federatedComputeEncryptionKeys) {
+                                LogUtil.d(
+                                        TAG,
+                                        "BackgroundKeyFetchJobService %d is done, fetched %d keys",
+                                        params.getJobId(),
+                                        federatedComputeEncryptionKeys.size());
+                                jobFinished(params, false/* wantsReschedule= */);
+                            }
+
+                            @Override
+                            public void onFailure(Throwable throwable) {
+                                LogUtil.e(
+                                        TAG,
+                                        "Failed to run job %d to fetch key and delete expired keys",
+                                        params.getJobId());
+                                if (throwable instanceof ExecutionException) {
+                                    LogUtil.e(
+                                            TAG,
+                                            "Background key fetch failed due to internal error");
+                                } else if (throwable instanceof TimeoutException) {
+                                    LogUtil.e(
+                                            TAG,
+                                            "Background key fetch failed due to timeout error");
+                                } else if (throwable instanceof InterruptedException) {
+                                    LogUtil.e(
+                                            TAG,
+                                            "Background key fetch failed due to interruption "
+                                            + "error");
+                                } else {
+                                    LogUtil.e(
+                                            TAG,
+                                            "Background key fetch failed due to unexpected error");
+                                }
+                                jobFinished(params, false /* wantsReschedule= */);
+                            }
+                        },
+                        mInjector.getLightWeightExecutor());
+        return true;
+    }
+
+    @Override
+    public boolean onStopJob(JobParameters params) {
+        LogUtil.d(TAG, "BackgroundKeyFetchJobService.onStopJob %d", params.getJobId());
+        return false;
+    }
+
+    /** Schedule the periodic background key fetch and delete job if it is not scheduled. */
+    public static boolean scheduleJobIfNeeded(Context context, Flags flags) {
+        final JobScheduler jobScheduler = context.getSystemService(JobScheduler.class);
+        if (jobScheduler == null) {
+            LogUtil.e(TAG, "Failed to get job scheduler from system service.");
+            return false;
+        }
+
+        final JobInfo scheduledJob = jobScheduler.getPendingJob(ENCRYPTION_KEY_FETCH_JOB_ID);
+        final JobInfo jobInfo =
+                new JobInfo.Builder(
+                                ENCRYPTION_KEY_FETCH_JOB_ID,
+                                new ComponentName(context, BackgroundKeyFetchJobService.class))
+                        .setPeriodic(
+                                flags.getEncryptionKeyFetchPeriodSeconds()
+                                        * 1000) // convert to milliseconds
+                        .setRequiresBatteryNotLow(true)
+                        .setRequiresDeviceIdle(true)
+                        .setRequiredNetworkType(JobInfo.NETWORK_TYPE_UNMETERED)
+                        .build();
+
+        if (!jobInfo.equals(scheduledJob)) {
+            jobScheduler.schedule(jobInfo);
+            LogUtil.d(
+                    TAG,
+                    "Scheduled job BackgroundKeyFetchJobService id %d",
+                    ENCRYPTION_KEY_FETCH_JOB_ID);
+            return true;
+        } else {
+            LogUtil.d(
+                    TAG,
+                    "Already scheduled job BackgroundKeyFetchJobService id %d",
+                    ENCRYPTION_KEY_FETCH_JOB_ID);
+            return false;
+        }
+    }
+}
diff --git a/federatedcompute/src/com/android/federatedcompute/services/encryption/FederatedComputeEncryptionKeyManager.java b/federatedcompute/src/com/android/federatedcompute/services/encryption/FederatedComputeEncryptionKeyManager.java
new file mode 100644
index 00000000..09df11c7
--- /dev/null
+++ b/federatedcompute/src/com/android/federatedcompute/services/encryption/FederatedComputeEncryptionKeyManager.java
@@ -0,0 +1,311 @@
+/*
+ * 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.encryption;
+
+import android.content.Context;
+
+import com.android.federatedcompute.internal.util.LogUtil;
+import com.android.federatedcompute.services.common.Clock;
+import com.android.federatedcompute.services.common.FederatedComputeExecutors;
+import com.android.federatedcompute.services.common.Flags;
+import com.android.federatedcompute.services.common.FlagsFactory;
+import com.android.federatedcompute.services.common.MonotonicClock;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKey;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKeyDao;
+import com.android.federatedcompute.services.http.FederatedComputeHttpRequest;
+import com.android.federatedcompute.services.http.FederatedComputeHttpResponse;
+import com.android.federatedcompute.services.http.HttpClient;
+import com.android.federatedcompute.services.http.HttpClientUtil;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.FluentFuture;
+
+import org.json.JSONArray;
+import org.json.JSONException;
+import org.json.JSONObject;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
+
+/** Class to manage key fetch. */
+public class FederatedComputeEncryptionKeyManager {
+    private static final String TAG = "FederatedComputeEncryptionKeyManager";
+
+    private interface EncryptionKeyResponseContract {
+        String RESPONSE_HEADER_CACHE_CONTROL_LABEL = "cache-control";
+        String RESPONSE_HEADER_AGE_LABEL = "age";
+
+        String RESPONSE_HEADER_CACHE_CONTROL_MAX_AGE_LABEL = "max-age=";
+
+        String RESPONSE_KEYS_LABEL = "keys";
+
+        String RESPONSE_KEY_ID_LABEL = "id";
+
+        String RESPONSE_PUBLIC_KEY = "key";
+    }
+
+    @VisibleForTesting private final FederatedComputeEncryptionKeyDao mEncryptionKeyDao;
+
+    private static volatile FederatedComputeEncryptionKeyManager sBackgroundKeyManager;
+
+    private final Clock mClock;
+
+    private final Flags mFlags;
+
+    private final HttpClient mHttpClient;
+
+    private final ExecutorService mBackgroundExecutor;
+
+    public FederatedComputeEncryptionKeyManager(
+            Clock clock,
+            FederatedComputeEncryptionKeyDao encryptionKeyDao,
+            Flags flags,
+            HttpClient httpClient,
+            ExecutorService backgroundExecutor) {
+        mClock = clock;
+        mEncryptionKeyDao = encryptionKeyDao;
+        mFlags = flags;
+        mHttpClient = httpClient;
+        mBackgroundExecutor = backgroundExecutor;
+    }
+
+    /**
+     * @return a singleton instance for key manager
+     */
+    public static FederatedComputeEncryptionKeyManager getInstance(Context context) {
+        if (sBackgroundKeyManager == null) {
+            synchronized (FederatedComputeEncryptionKeyManager.class) {
+                if (sBackgroundKeyManager == null) {
+                    FederatedComputeEncryptionKeyDao encryptionKeyDao =
+                            FederatedComputeEncryptionKeyDao.getInstance(context);
+                    HttpClient client = new HttpClient();
+                    Clock clock = MonotonicClock.getInstance();
+                    Flags flags = FlagsFactory.getFlags();
+                    sBackgroundKeyManager =
+                            new FederatedComputeEncryptionKeyManager(
+                                    clock,
+                                    encryptionKeyDao,
+                                    flags,
+                                    client,
+                                    FederatedComputeExecutors.getBackgroundExecutor());
+                }
+            }
+        }
+        return sBackgroundKeyManager;
+    }
+
+    /** For testing only, returns an instance of key manager for test. */
+    @VisibleForTesting
+    public static FederatedComputeEncryptionKeyManager getInstanceForTest(
+            Clock clock,
+            FederatedComputeEncryptionKeyDao encryptionKeyDao,
+            Flags flags,
+            HttpClient client,
+            ExecutorService executor) {
+        if (sBackgroundKeyManager == null) {
+            synchronized (FederatedComputeEncryptionKeyManager.class) {
+                if (sBackgroundKeyManager == null) {
+                    sBackgroundKeyManager =
+                            new FederatedComputeEncryptionKeyManager(
+                                    clock, encryptionKeyDao, flags, client, executor);
+                }
+            }
+        }
+        return sBackgroundKeyManager;
+    }
+
+    /**
+     * Fetch the active key from the server, persists the fetched key to encryption_key table, and
+     * deletes expired keys
+     */
+    public FluentFuture<List<FederatedComputeEncryptionKey>> fetchAndPersistActiveKeys(
+            @FederatedComputeEncryptionKey.KeyType int keyType, boolean isScheduledJob) {
+        String fetchUri = mFlags.getEncryptionKeyFetchUrl();
+        if (fetchUri == null) {
+            throw new IllegalArgumentException("Url to fetch active encryption keys is null");
+        }
+
+        FederatedComputeHttpRequest request =
+                FederatedComputeHttpRequest.create(
+                        fetchUri,
+                        HttpClientUtil.HttpMethod.GET,
+                        new HashMap<String, String>(),
+                        HttpClientUtil.EMPTY_BODY);
+
+        return FluentFuture.from(mHttpClient.performRequestAsyncWithRetry(request))
+                .transform(
+                        response ->
+                                parseFetchEncryptionKeyPayload(
+                                        response, keyType, mClock.currentTimeMillis()),
+                        mBackgroundExecutor)
+                .transform(
+                        result -> {
+                            result.forEach(mEncryptionKeyDao::insertEncryptionKey);
+                            if (isScheduledJob) {
+                                // When the job is a background scheduled job, delete the expired
+                                // keys, otherwise, only fetch from the key server.
+                                mEncryptionKeyDao.deleteExpiredKeys();
+                            }
+                            return result;
+                        },
+                        mBackgroundExecutor); // TODO: Add timeout controlled by Ph flags
+    }
+
+    private ImmutableList<FederatedComputeEncryptionKey> parseFetchEncryptionKeyPayload(
+            FederatedComputeHttpResponse keyFetchResponse,
+            @FederatedComputeEncryptionKey.KeyType int keyType,
+            Long fetchTime) {
+        String payload = new String(Objects.requireNonNull(keyFetchResponse.getPayload()));
+        Map<String, List<String>> headers = keyFetchResponse.getHeaders();
+        long ttlInSeconds = getTTL(headers);
+        if (ttlInSeconds <= 0) {
+            ttlInSeconds = mFlags.getFederatedComputeEncryptionKeyMaxAgeSeconds();
+        }
+
+        try {
+            JSONObject responseObj = new JSONObject(payload);
+            JSONArray keysArr =
+                    responseObj.getJSONArray(EncryptionKeyResponseContract.RESPONSE_KEYS_LABEL);
+            ImmutableList.Builder<FederatedComputeEncryptionKey> encryptionKeys =
+                    ImmutableList.builder();
+
+            for (int i = 0; i < keysArr.length(); i++) {
+                JSONObject keyObj = keysArr.getJSONObject(i);
+                FederatedComputeEncryptionKey key =
+                        new FederatedComputeEncryptionKey.Builder()
+                                .setKeyIdentifier(
+                                        keyObj.getString(
+                                                EncryptionKeyResponseContract
+                                                        .RESPONSE_KEY_ID_LABEL))
+                                .setPublicKey(
+                                        keyObj.getString(
+                                                EncryptionKeyResponseContract.RESPONSE_PUBLIC_KEY))
+                                .setKeyType(keyType)
+                                .setCreationTime(fetchTime)
+                                .setExpiryTime(
+                                        fetchTime + ttlInSeconds * 1000) // convert to milliseconds
+                                .build();
+                encryptionKeys.add(key);
+            }
+            return encryptionKeys.build();
+        } catch (JSONException e) {
+            LogUtil.e(TAG, "Invalid Json response: " + e.getMessage());
+            return ImmutableList.of();
+        }
+    }
+
+    /**
+     * Parse the "age" and "cache-control" of response headers. Calculate the ttl of the current key
+     * maxage (in cache-control) - age.
+     *
+     * @return the ttl in seconds of the keys.
+     */
+    @VisibleForTesting
+    static long getTTL(Map<String, List<String>> headers) {
+        String cacheControl = null;
+        int cachedAge = 0;
+        int remainingHeaders = 2;
+        for (String key : headers.keySet()) {
+            if (key != null) {
+                if (key.equalsIgnoreCase(
+                        EncryptionKeyResponseContract.RESPONSE_HEADER_CACHE_CONTROL_LABEL)) {
+                    List<String> field = headers.get(key);
+                    if (field != null && field.size() > 0) {
+                        cacheControl = field.get(0).toLowerCase(Locale.ENGLISH);
+                        remainingHeaders -= 1;
+                    }
+
+                } else if (key.equalsIgnoreCase(
+                        EncryptionKeyResponseContract.RESPONSE_HEADER_AGE_LABEL)) {
+                    List<String> field = headers.get(key);
+                    if (field != null && field.size() > 0) {
+                        try {
+                            cachedAge = Integer.parseInt(field.get(0));
+                        } catch (NumberFormatException e) {
+                            LogUtil.e(TAG, "Error parsing age header");
+                        }
+                        remainingHeaders -= 1;
+                    }
+                }
+            }
+            if (remainingHeaders == 0) {
+                break;
+            }
+        }
+        if (cacheControl == null) {
+            LogUtil.d(TAG, "Cache-Control header or value is missing");
+            return 0;
+        }
+
+        String[] tokens = cacheControl.split(",", /* limit= */ 0);
+        long maxAge = 0;
+        for (String s : tokens) {
+            String token = s.trim();
+            if (token.startsWith(
+                    EncryptionKeyResponseContract.RESPONSE_HEADER_CACHE_CONTROL_MAX_AGE_LABEL)) {
+                try {
+                    maxAge =
+                            Long.parseLong(
+                                    token.substring(
+                                            /* beginIndex= */ EncryptionKeyResponseContract
+                                                    .RESPONSE_HEADER_CACHE_CONTROL_MAX_AGE_LABEL
+                                                    .length())); // in the format of
+                    // "max-age=<number>"
+                } catch (NumberFormatException e) {
+                    LogUtil.d(TAG, "Failed to parse max-age value");
+                    return 0;
+                }
+            }
+        }
+        if (maxAge == 0) {
+            LogUtil.d(TAG, "max-age directive is missing");
+            return 0;
+        }
+        return maxAge - cachedAge;
+    }
+
+    /** Get active keys, if there is no active key, then force a fetch from the key service.
+     * In the case of key fetching from the key service, the http call
+     * is executed on a BlockingExecutor.
+     * @return The list of active keys.
+     */
+    public List<FederatedComputeEncryptionKey> getOrFetchActiveKeys(int keyType, int keyCount) {
+        List<FederatedComputeEncryptionKey> activeKeys = mEncryptionKeyDao
+                .getLatestExpiryNKeys(keyCount);
+        if (activeKeys.size() > 0) {
+            return activeKeys;
+        }
+        try {
+            var fetchedKeysUnused = fetchAndPersistActiveKeys(keyType,
+                    /* isScheduledJob= */ false).get(1, TimeUnit.SECONDS);
+            activeKeys = mEncryptionKeyDao.getLatestExpiryNKeys(keyCount);
+            if (activeKeys.size() > 0) {
+                return activeKeys;
+            }
+        } catch (Exception e) {
+            LogUtil.e(TAG, "Exception encountered when forcing encryption key fetch: "
+                    + e.getMessage());
+        }
+        return activeKeys;
+    }
+}
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/CheckinResult.java b/federatedcompute/src/com/android/federatedcompute/services/http/CheckinResult.java
index e82bfe61..00542f4e 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/CheckinResult.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/CheckinResult.java
@@ -21,6 +21,7 @@ import android.annotation.Nullable;
 import com.android.internal.util.Preconditions;
 
 import com.google.internal.federated.plan.ClientOnlyPlan;
+import com.google.internal.federatedcompute.v1.RejectionInfo;
 import com.google.ondevicepersonalization.federatedcompute.proto.TaskAssignment;
 
 /**
@@ -30,7 +31,7 @@ public class CheckinResult {
     private String mInputCheckpoint = null;
     private ClientOnlyPlan mPlanData = null;
     private TaskAssignment mTaskAssignment = null;
-
+    private RejectionInfo mRejectionInfo = null;
     public CheckinResult(
             String inputCheckpoint, ClientOnlyPlan planData, TaskAssignment taskAssignment) {
         this.mInputCheckpoint = inputCheckpoint;
@@ -38,6 +39,10 @@ public class CheckinResult {
         this.mTaskAssignment = taskAssignment;
     }
 
+    public CheckinResult(RejectionInfo mRejectionInfo) {
+        this.mRejectionInfo = mRejectionInfo;
+    }
+
     @Nullable
     public String getInputCheckpointFile() {
         Preconditions.checkArgument(
@@ -55,4 +60,9 @@ public class CheckinResult {
     public TaskAssignment getTaskAssignment() {
         return mTaskAssignment;
     }
+
+    @Nullable
+    public RejectionInfo getRejectionInfo() {
+        return mRejectionInfo;
+    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequest.java b/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequest.java
index 527811df..80ef1391 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequest.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequest.java
@@ -16,9 +16,7 @@
 
 package com.android.federatedcompute.services.http;
 
-import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_ENCODING_HDR;
 import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_LENGTH_HDR;
-import static com.android.federatedcompute.services.http.HttpClientUtil.GZIP_ENCODING_HDR;
 
 import com.android.federatedcompute.services.http.HttpClientUtil.HttpMethod;
 
@@ -45,18 +43,10 @@ public final class FederatedComputeHttpRequest {
 
     /** Creates a {@link FederatedComputeHttpRequest} based on given inputs. */
     public static FederatedComputeHttpRequest create(
-            String uri,
-            HttpMethod httpMethod,
-            HashMap<String, String> extraHeaders,
-            byte[] body,
-            boolean useCompression) {
+            String uri, HttpMethod httpMethod, HashMap<String, String> extraHeaders, byte[] body) {
         if (!uri.startsWith(HTTPS_SCHEMA) && !uri.startsWith(LOCAL_HOST_URI)) {
             throw new IllegalArgumentException("Non-HTTPS URIs are not supported: " + uri);
         }
-        if (useCompression) {
-            body = HttpClientUtil.compressWithGzip(body);
-            extraHeaders.put(CONTENT_ENCODING_HDR, GZIP_ENCODING_HDR);
-        }
         if (extraHeaders.containsKey(CONTENT_LENGTH_HDR)) {
             throw new IllegalArgumentException("Content-Length header should not be provided!");
         }
@@ -67,7 +57,6 @@ public final class FederatedComputeHttpRequest {
             }
             extraHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(body.length));
         }
-
         return new FederatedComputeHttpRequest(uri, httpMethod, extraHeaders, body);
     }
 
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponse.java b/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponse.java
index 5eb1e920..48338f29 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponse.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponse.java
@@ -16,16 +16,20 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_ENCODING_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.GZIP_ENCODING_HDR;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 /** Class to hold FederatedCompute http response. */
 public class FederatedComputeHttpResponse {
     private Integer mStatusCode;
-    private Map<String, List<String>> mHeaders;
+    private Map<String, List<String>> mHeaders = new HashMap<>();
     private byte[] mPayload;
 
     private FederatedComputeHttpResponse() {}
@@ -45,6 +49,18 @@ public class FederatedComputeHttpResponse {
         return mPayload;
     }
 
+    /** Returns whether http response body is compressed with gzip. */
+    public boolean isResponseCompressed() {
+        if (mHeaders.containsKey(CONTENT_ENCODING_HDR)) {
+            for (String format : mHeaders.get(CONTENT_ENCODING_HDR)) {
+                if (format.contains(GZIP_ENCODING_HDR)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     /** Builder for FederatedComputeHttpResponse. */
     public static final class Builder {
         private final FederatedComputeHttpResponse mHttpResponse;
@@ -71,6 +87,7 @@ public class FederatedComputeHttpResponse {
             mHttpResponse.mPayload = payload;
             return this;
         }
+
         /** Build {@link FederatedComputeHttpResponse}. */
         public FederatedComputeHttpResponse build() {
             if (mHttpResponse.mStatusCode == null) {
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/HttpClient.java b/federatedcompute/src/com/android/federatedcompute/services/http/HttpClient.java
index 2caeccea..acbd8c0e 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/HttpClient.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/HttpClient.java
@@ -45,6 +45,7 @@ import java.util.concurrent.TimeUnit;
  */
 public class HttpClient {
     private static final String TAG = HttpClient.class.getSimpleName();
+    private static final int RETRY_LIMIT = 3;
     private static final int NETWORK_CONNECT_TIMEOUT_MS = (int) TimeUnit.SECONDS.toMillis(5);
     private static final int NETWORK_READ_TIMEOUT_MS = (int) TimeUnit.SECONDS.toMillis(30);
 
@@ -60,17 +61,45 @@ public class HttpClient {
         return urlConnection;
     }
 
-    /** Perform HTTP requests based on given information asynchronously. */
+    /**
+     * Perform HTTP requests based on given information asynchronously with retries in case http
+     * will return not OK response code.
+     */
     @NonNull
-    public ListenableFuture<FederatedComputeHttpResponse> performRequestAsync(
+    public ListenableFuture<FederatedComputeHttpResponse> performRequestAsyncWithRetry(
             FederatedComputeHttpRequest request) {
         try {
-            return getBlockingExecutor().submit(() -> performRequest(request));
+            return getBlockingExecutor().submit(() -> performRequestWithRetry(request));
         } catch (Exception e) {
             return Futures.immediateFailedFuture(e);
         }
     }
 
+    /** Perform HTTP requests based on given information with retries. */
+    @NonNull
+    public FederatedComputeHttpResponse performRequestWithRetry(FederatedComputeHttpRequest request)
+            throws IOException {
+        int count = 0;
+        FederatedComputeHttpResponse response = null;
+        while (count < RETRY_LIMIT) {
+            try {
+                response = performRequest(request);
+                if (HTTP_OK_STATUS.contains(response.getStatusCode())) {
+                    return response;
+                }
+                // we want to continue retry in case it is IO exception.
+            } catch (IOException e) {
+                // propagate IO exception after RETRY_LIMIT times attempt.
+                if (count >= RETRY_LIMIT - 1) {
+                    throw e;
+                }
+            } finally {
+                count++;
+            }
+        }
+        return response;
+    }
+
     /** Perform HTTP requests based on given information. */
     @NonNull
     public FederatedComputeHttpResponse performRequest(FederatedComputeHttpRequest request)
@@ -93,7 +122,7 @@ public class HttpClient {
             urlConnection = (HttpURLConnection) setup(url);
         } catch (IOException e) {
             LogUtil.e(TAG, e, "Failed to open target URL");
-            throw new IllegalArgumentException("Failed to open target URL", e);
+            throw new IOException("Failed to open target URL", e);
         }
 
         try {
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/HttpClientUtil.java b/federatedcompute/src/com/android/federatedcompute/services/http/HttpClientUtil.java
index 1d6bf784..61d719e9 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/HttpClientUtil.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/HttpClientUtil.java
@@ -21,23 +21,25 @@ import com.android.federatedcompute.internal.util.LogUtil;
 import com.google.common.collect.ImmutableSet;
 import com.google.protobuf.ByteString;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
 
 /** Utility class containing http related variable e.g. headers, method. */
 public final class HttpClientUtil {
     private static final String TAG = HttpClientUtil.class.getSimpleName();
-    public static final String IDENTITY_ENCODING_HDR = "identity";
     public static final String CONTENT_ENCODING_HDR = "Content-Encoding";
+
+    public static final String ACCEPT_ENCODING_HDR = "Accept-Encoding";
     public static final String CONTENT_LENGTH_HDR = "Content-Length";
     public static final String GZIP_ENCODING_HDR = "gzip";
-    public static final String API_KEY_HDR = "x-goog-api-key";
     public static final String CONTENT_TYPE_HDR = "Content-Type";
     public static final String PROTOBUF_CONTENT_TYPE = "application/x-protobuf";
     public static final String OCTET_STREAM = "application/octet-stream";
-    public static final String CLIENT_DECODE_GZIP_SUFFIX = "+gzip";
     public static final ImmutableSet<Integer> HTTP_OK_STATUS = ImmutableSet.of(200, 201);
-    public static final String FAKE_API_KEY = "FAKE_API_KEY";
+    public static final String ODP_IDEMPOTENCY_KEY = "odp-idempotency-key";
     public static final int DEFAULT_BUFFER_SIZE = 1024;
     public static final byte[] EMPTY_BODY = new byte[0];
 
@@ -57,7 +59,24 @@ public final class HttpClientUtil {
             return outputStream.toByteString().toByteArray();
         } catch (IOException e) {
             LogUtil.e(TAG, "Failed to compress using Gzip");
-            throw new IllegalArgumentException("Failed to compress using Gzip", e);
+            throw new IllegalStateException("Failed to compress using Gzip", e);
+        }
+    }
+
+    /** Uncompresses the input data using Gzip. */
+    public static byte[] uncompressWithGzip(byte[] data) {
+        try (ByteArrayInputStream inputStream = new ByteArrayInputStream(data);
+                GZIPInputStream gzip = new GZIPInputStream(inputStream);
+                ByteArrayOutputStream result = new ByteArrayOutputStream()) {
+            int length;
+            byte[] buffer = new byte[DEFAULT_BUFFER_SIZE];
+            while ((length = gzip.read(buffer, 0, DEFAULT_BUFFER_SIZE)) > 0) {
+                result.write(buffer, 0, length);
+            }
+            return result.toByteArray();
+        } catch (Exception e) {
+            LogUtil.e(TAG, "Failed to decompress the data.", e);
+            throw new IllegalStateException("Failed to unscompress using Gzip", e);
         }
     }
 
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java b/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
index a55fcf49..1bae7ebd 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/HttpFederatedProtocol.java
@@ -16,12 +16,21 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_HTTP_ISSUE_CHECKIN;
+import static com.android.federatedcompute.services.common.Constants.TRACE_HTTP_REPORT_RESULT;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getBackgroundExecutor;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getLightweightExecutor;
 import static com.android.federatedcompute.services.common.FileUtils.createTempFile;
 import static com.android.federatedcompute.services.common.FileUtils.readFileAsByteArray;
 import static com.android.federatedcompute.services.common.FileUtils.writeToFile;
+import static com.android.federatedcompute.services.http.HttpClientUtil.ACCEPT_ENCODING_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.GZIP_ENCODING_HDR;
 import static com.android.federatedcompute.services.http.HttpClientUtil.HTTP_OK_STATUS;
+import static com.android.federatedcompute.services.http.HttpClientUtil.ODP_IDEMPOTENCY_KEY;
+import static com.android.federatedcompute.services.http.HttpClientUtil.compressWithGzip;
+import static com.android.federatedcompute.services.http.HttpClientUtil.uncompressWithGzip;
+
+import android.os.Trace;
 
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.http.HttpClientUtil.HttpMethod;
@@ -29,13 +38,15 @@ import com.android.federatedcompute.services.training.util.ComputationResult;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
-import com.google.common.util.concurrent.AsyncCallable;
 import com.google.common.util.concurrent.FluentFuture;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.internal.federated.plan.ClientOnlyPlan;
 import com.google.internal.federatedcompute.v1.ClientVersion;
+import com.google.internal.federatedcompute.v1.RejectionInfo;
 import com.google.internal.federatedcompute.v1.Resource;
+import com.google.internal.federatedcompute.v1.ResourceCapabilities;
+import com.google.internal.federatedcompute.v1.ResourceCompressionFormat;
 import com.google.ondevicepersonalization.federatedcompute.proto.CreateTaskAssignmentRequest;
 import com.google.ondevicepersonalization.federatedcompute.proto.CreateTaskAssignmentResponse;
 import com.google.ondevicepersonalization.federatedcompute.proto.ReportResultRequest;
@@ -46,10 +57,12 @@ import com.google.ondevicepersonalization.federatedcompute.proto.UploadInstructi
 import com.google.protobuf.InvalidProtocolBufferException;
 
 import java.util.HashMap;
+import java.util.UUID;
+import java.util.concurrent.Callable;
 
 /** Implements a single session of HTTP-based federated compute protocol. */
 public final class HttpFederatedProtocol {
-    public static final String TAG = "HttpFederatedProtocol";
+    public static final String TAG = HttpFederatedProtocol.class.getSimpleName();
     private final String mClientVersion;
     private final String mPopulationName;
     private final HttpClient mHttpClient;
@@ -64,8 +77,7 @@ public final class HttpFederatedProtocol {
         this.mClientVersion = clientVersion;
         this.mPopulationName = populationName;
         this.mHttpClient = httpClient;
-        this.mTaskAssignmentRequestCreator =
-                new ProtocolRequestCreator(entryUri, new HashMap<>(), false);
+        this.mTaskAssignmentRequestCreator = new ProtocolRequestCreator(entryUri, new HashMap<>());
     }
 
     /** Creates a HttpFederatedProtocol object. */
@@ -76,53 +88,77 @@ public final class HttpFederatedProtocol {
 
     /** Helper function to perform check in and download federated task from remote servers. */
     public ListenableFuture<CheckinResult> issueCheckin() {
-        ListenableFuture<TaskAssignment> taskAssignmentFuture =
-                FluentFuture.from(createTaskAssignment())
-                        .transform(
-                                getTaskAssignmentHttpResponse ->
-                                        getTaskAssignment(getTaskAssignmentHttpResponse),
-                                getLightweightExecutor());
-        ListenableFuture<FederatedComputeHttpResponse> planDataResponseFuture =
-                FluentFuture.from(taskAssignmentFuture)
-                        .transformAsync(
-                                taskAssignment -> fetchTaskResource(taskAssignment.getPlan()),
-                                getBackgroundExecutor());
-        ListenableFuture<FederatedComputeHttpResponse> checkpointDataResponseFuture =
-                FluentFuture.from(taskAssignmentFuture)
-                        .transformAsync(
-                                taskAssignment ->
-                                        fetchTaskResource(taskAssignment.getInitCheckpoint()),
-                                getBackgroundExecutor());
-        return Futures.whenAllSucceed(
-                        taskAssignmentFuture, planDataResponseFuture, checkpointDataResponseFuture)
-                .callAsync(
-                        new AsyncCallable<CheckinResult>() {
-                            @Override
-                            public ListenableFuture<CheckinResult> call() {
-                                return getCheckinResult(
-                                        planDataResponseFuture,
-                                        checkpointDataResponseFuture,
-                                        taskAssignmentFuture);
+        Trace.beginAsyncSection(TRACE_HTTP_ISSUE_CHECKIN, 0);
+        return FluentFuture.from(createTaskAssignment())
+                .transformAsync(
+                        federatedComputeHttpResponse -> {
+                            validateHttpResponseStatus(
+                                    "Start task assignment", federatedComputeHttpResponse);
+                            CreateTaskAssignmentResponse taskAssignmentResponse;
+                            try {
+                                taskAssignmentResponse =
+                                        CreateTaskAssignmentResponse.parseFrom(
+                                                federatedComputeHttpResponse.getPayload());
+                            } catch (InvalidProtocolBufferException e) {
+                                throw new IllegalStateException(
+                                        "Could not parse StartTaskAssignmentResponse proto", e);
+                            }
+                            if (taskAssignmentResponse.hasRejectionInfo()) {
+                                return Futures.immediateFuture(
+                                        new CheckinResult(
+                                                taskAssignmentResponse.getRejectionInfo()));
                             }
+                            TaskAssignment taskAssignment =
+                                    getTaskAssignment(taskAssignmentResponse);
+                            ListenableFuture<FederatedComputeHttpResponse> planDataResponseFuture =
+                                    fetchTaskResource(taskAssignment.getPlan());
+                            ListenableFuture<FederatedComputeHttpResponse>
+                                    checkpointDataResponseFuture =
+                                            fetchTaskResource(taskAssignment.getInitCheckpoint());
+                            return Futures.whenAllSucceed(
+                                            planDataResponseFuture, checkpointDataResponseFuture)
+                                    .call(
+                                            new Callable<CheckinResult>() {
+                                                @Override
+                                                public CheckinResult call() throws Exception {
+                                                    return getCheckinResult(
+                                                            planDataResponseFuture,
+                                                            checkpointDataResponseFuture,
+                                                            taskAssignment);
+                                                }
+                                            },
+                                            getBackgroundExecutor());
                         },
                         getBackgroundExecutor());
     }
 
     /** Helper functions to reporting result and upload result. */
-    public FluentFuture<Void> reportResult(ComputationResult computationResult) {
+    public FluentFuture<RejectionInfo> reportResult(ComputationResult computationResult) {
+        Trace.beginAsyncSection(TRACE_HTTP_REPORT_RESULT, 0);
         if (computationResult != null && computationResult.isResultSuccess()) {
             return FluentFuture.from(performReportResult(computationResult))
                     .transformAsync(
-                            reportResp ->
-                                    processReportResultResponseAndUploadResult(
-                                            reportResp, computationResult),
-                            getBackgroundExecutor())
-                    .transform(
-                            resp -> {
-                                validateHttpResponseStatus("Upload result", resp);
-                                return null;
+                            reportResp -> {
+                                ReportResultResponse reportResultResponse =
+                                        getReportResultResponse(reportResp);
+                                if (reportResultResponse.hasRejectionInfo()) {
+                                    return Futures.immediateFuture(
+                                            reportResultResponse.getRejectionInfo());
+                                }
+                                return FluentFuture.from(
+                                                processReportResultResponseAndUploadResult(
+                                                        reportResultResponse, computationResult))
+                                        .transform(
+                                                resp -> {
+                                                    validateHttpResponseStatus(
+                                                            "Upload result", resp);
+                                                    Trace.endAsyncSection(
+                                                            TRACE_HTTP_REPORT_RESULT, 0);
+                                                    return null;
+                                                },
+                                                getLightweightExecutor());
                             },
-                            getLightweightExecutor());
+                            getBackgroundExecutor());
         } else {
             return FluentFuture.from(performReportResult(computationResult))
                     .transform(
@@ -138,7 +174,13 @@ public final class HttpFederatedProtocol {
         CreateTaskAssignmentRequest request =
                 CreateTaskAssignmentRequest.newBuilder()
                         .setClientVersion(ClientVersion.newBuilder().setVersionCode(mClientVersion))
+                        .setResourceCapabilities(
+                                ResourceCapabilities.newBuilder()
+                                        .addSupportedCompressionFormats(
+                                                ResourceCompressionFormat
+                                                        .RESOURCE_COMPRESSION_FORMAT_GZIP))
                         .build();
+
         String taskAssignmentUriSuffix =
                 String.format(
                         "/taskassignment/v1/population/%1$s:create-task-assignment",
@@ -147,20 +189,16 @@ public final class HttpFederatedProtocol {
                 mTaskAssignmentRequestCreator.createProtoRequest(
                         taskAssignmentUriSuffix,
                         HttpMethod.POST,
+                        new HashMap<>(),
                         request.toByteArray(),
                         /* isProtobufEncoded= */ true);
-        return mHttpClient.performRequestAsync(httpRequest);
+        httpRequest
+                .getExtraHeaders()
+                .put(ODP_IDEMPOTENCY_KEY, System.currentTimeMillis() + " - " + UUID.randomUUID());
+        return mHttpClient.performRequestAsyncWithRetry(httpRequest);
     }
 
-    private TaskAssignment getTaskAssignment(FederatedComputeHttpResponse httpResponse) {
-        validateHttpResponseStatus("Start task assignment", httpResponse);
-        CreateTaskAssignmentResponse taskAssignmentResponse;
-        try {
-            taskAssignmentResponse =
-                    CreateTaskAssignmentResponse.parseFrom(httpResponse.getPayload());
-        } catch (InvalidProtocolBufferException e) {
-            throw new IllegalStateException("Could not parse StartTaskAssignmentResponse proto", e);
-        }
+    private TaskAssignment getTaskAssignment(CreateTaskAssignmentResponse taskAssignmentResponse) {
         if (taskAssignmentResponse.hasRejectionInfo()) {
             throw new IllegalStateException("Device rejected by server.");
         }
@@ -194,32 +232,44 @@ public final class HttpFederatedProtocol {
         this.mAssignmentId = taskAssignment.getAssignmentId();
     }
 
-    private ListenableFuture<CheckinResult> getCheckinResult(
+    private CheckinResult getCheckinResult(
             ListenableFuture<FederatedComputeHttpResponse> planDataResponseFuture,
             ListenableFuture<FederatedComputeHttpResponse> checkpointDataResponseFuture,
-            ListenableFuture<TaskAssignment> taskAssignmentFuture) {
+            TaskAssignment taskAssignment)
+            throws Exception {
+
+        FederatedComputeHttpResponse planDataResponse = Futures.getDone(planDataResponseFuture);
+        FederatedComputeHttpResponse checkpointDataResponse =
+                Futures.getDone(checkpointDataResponseFuture);
+        validateHttpResponseStatus("Fetch plan", planDataResponse);
+        validateHttpResponseStatus("Fetch checkpoint", checkpointDataResponse);
+
+        // Process download ClientOnlyPlan.
+        byte[] planData = planDataResponse.getPayload();
+        if (taskAssignment.getPlan().getCompressionFormat()
+                        == ResourceCompressionFormat.RESOURCE_COMPRESSION_FORMAT_GZIP
+                || planDataResponse.isResponseCompressed()) {
+            planData = uncompressWithGzip(planData);
+        }
+        ClientOnlyPlan clientOnlyPlan;
         try {
-            FederatedComputeHttpResponse planDataResponse = Futures.getDone(planDataResponseFuture);
-            FederatedComputeHttpResponse checkpointDataResponse =
-                    Futures.getDone(checkpointDataResponseFuture);
-            TaskAssignment taskAssignment = Futures.getDone(taskAssignmentFuture);
-            validateHttpResponseStatus("Fetch plan", planDataResponse);
-            validateHttpResponseStatus("Fetch checkpoint", checkpointDataResponse);
-            ClientOnlyPlan clientOnlyPlan;
-            try {
-                clientOnlyPlan = ClientOnlyPlan.parseFrom(planDataResponse.getPayload());
-            } catch (InvalidProtocolBufferException e) {
-                LogUtil.e(TAG, e, "Could not parse ClientOnlyPlan proto");
-                return Futures.immediateFailedFuture(
-                        new IllegalStateException("Could not parse ClientOnlyPlan proto", e));
-            }
-            String inputCheckpointFile = createTempFile("input", ".ckp");
-            writeToFile(inputCheckpointFile, checkpointDataResponse.getPayload());
-            return Futures.immediateFuture(
-                    new CheckinResult(inputCheckpointFile, clientOnlyPlan, taskAssignment));
-        } catch (Exception e) {
-            return Futures.immediateFailedFuture(e);
+            clientOnlyPlan = ClientOnlyPlan.parseFrom(planData);
+        } catch (InvalidProtocolBufferException e) {
+            LogUtil.e(TAG, e, "Could not parse ClientOnlyPlan proto");
+            throw new IllegalStateException("Could not parse ClientOnlyPlan proto", e);
         }
+
+        // Process download checkpoint resource.
+        String inputCheckpointFile = createTempFile("input", ".ckp");
+        byte[] checkpointData = checkpointDataResponse.getPayload();
+        if (taskAssignment.getInitCheckpoint().getCompressionFormat()
+                        == ResourceCompressionFormat.RESOURCE_COMPRESSION_FORMAT_GZIP
+                || checkpointDataResponse.isResponseCompressed()) {
+            checkpointData = uncompressWithGzip(checkpointData);
+        }
+        writeToFile(inputCheckpointFile, checkpointData);
+        Trace.endAsyncSection(TRACE_HTTP_ISSUE_CHECKIN, 0);
+        return new CheckinResult(inputCheckpointFile, clientOnlyPlan, taskAssignment);
     }
 
     private ListenableFuture<FederatedComputeHttpResponse> performReportResult(
@@ -249,31 +299,26 @@ public final class HttpFederatedProtocol {
                         HttpMethod.PUT,
                         startDataUploadRequest.toByteArray(),
                         /* isProtobufEncoded= */ true);
-        return mHttpClient.performRequestAsync(httpRequest);
+        return mHttpClient.performRequestAsyncWithRetry(httpRequest);
     }
 
     private ListenableFuture<FederatedComputeHttpResponse>
             processReportResultResponseAndUploadResult(
-                    FederatedComputeHttpResponse httpResponse,
+            ReportResultResponse reportResultResponse,
                     ComputationResult computationResult) {
         try {
-            validateHttpResponseStatus("ReportResult", httpResponse);
-            ReportResultResponse reportResultResponse =
-                    ReportResultResponse.parseFrom(httpResponse.getPayload());
-            // TODO(b/297605806): better handle rejection info.
-            if (reportResultResponse.hasRejectionInfo()) {
-                return Futures.immediateFailedFuture(
-                        new IllegalStateException(
-                                "ReportResult got rejection: " + httpResponse.getStatusCode()));
-            }
             Preconditions.checkArgument(
                     !computationResult.getOutputCheckpointFile().isEmpty(),
                     "Output checkpoint file should not be empty");
-            byte[] outputBytes = readFileAsByteArray(computationResult.getOutputCheckpointFile());
             UploadInstruction uploadInstruction = reportResultResponse.getUploadInstruction();
             Preconditions.checkArgument(
                     !uploadInstruction.getUploadLocation().isEmpty(),
                     "UploadInstruction.upload_location must not be empty");
+            byte[] outputBytes = readFileAsByteArray(computationResult.getOutputCheckpointFile());
+            if (uploadInstruction.getCompressionFormat()
+                    == ResourceCompressionFormat.RESOURCE_COMPRESSION_FORMAT_GZIP) {
+                outputBytes = compressWithGzip(outputBytes);
+            }
             HashMap<String, String> requestHeader = new HashMap<>();
             uploadInstruction
                     .getExtraRequestHeadersMap()
@@ -293,14 +338,21 @@ public final class HttpFederatedProtocol {
                             uploadInstruction.getUploadLocation(),
                             HttpMethod.PUT,
                             requestHeader,
-                            outputBytes,
-                            /* useCompression= */ false);
-            return mHttpClient.performRequestAsync(httpUploadRequest);
+                            outputBytes);
+            return mHttpClient.performRequestAsyncWithRetry(httpUploadRequest);
         } catch (Exception e) {
             return Futures.immediateFailedFuture(e);
         }
     }
 
+    private ReportResultResponse getReportResultResponse(FederatedComputeHttpResponse httpResponse)
+            throws InvalidProtocolBufferException {
+        validateHttpResponseStatus("ReportResult", httpResponse);
+        ReportResultResponse reportResultResponse =
+                ReportResultResponse.parseFrom(httpResponse.getPayload());
+        return reportResultResponse;
+    }
+
     private void validateHttpResponseStatus(
             String stage, FederatedComputeHttpResponse httpResponse) {
         if (!HTTP_OK_STATUS.contains(httpResponse.getStatusCode())) {
@@ -315,14 +367,22 @@ public final class HttpFederatedProtocol {
             case URI:
                 Preconditions.checkArgument(
                         !resource.getUri().isEmpty(), "Resource.uri must be non-empty when set");
+                HashMap<String, String> headerList = new HashMap<>();
+                if (resource.getCompressionFormat()
+                        == ResourceCompressionFormat.RESOURCE_COMPRESSION_FORMAT_GZIP) {
+                    // Set this header to disable decompressive transcoding when download from
+                    // Google Cloud Storage.
+                    // https://cloud.google.com/storage/docs/transcoding#decompressive_transcoding
+                    headerList.put(ACCEPT_ENCODING_HDR, GZIP_ENCODING_HDR);
+                }
+                LogUtil.d(TAG, "start fetch task resources");
                 FederatedComputeHttpRequest httpRequest =
                         FederatedComputeHttpRequest.create(
                                 resource.getUri(),
                                 HttpMethod.GET,
-                                new HashMap<String, String>(),
-                                HttpClientUtil.EMPTY_BODY,
-                                /* useCompression= */ false);
-                return mHttpClient.performRequestAsync(httpRequest);
+                                headerList,
+                                HttpClientUtil.EMPTY_BODY);
+                return mHttpClient.performRequestAsyncWithRetry(httpRequest);
             case INLINE_RESOURCE:
                 return Futures.immediateFailedFuture(
                         new UnsupportedOperationException("Inline resource is not supported yet."));
diff --git a/federatedcompute/src/com/android/federatedcompute/services/http/ProtocolRequestCreator.java b/federatedcompute/src/com/android/federatedcompute/services/http/ProtocolRequestCreator.java
index e87447f4..d73f1d05 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/http/ProtocolRequestCreator.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/http/ProtocolRequestCreator.java
@@ -32,28 +32,23 @@ import java.util.HashMap;
 public final class ProtocolRequestCreator {
     private final String mRequestBaseUri;
     private final HashMap<String, String> mHeaderList;
-    private boolean mUseCompression;
 
-    public ProtocolRequestCreator(
-            String requestBaseUri, HashMap<String, String> headerList, boolean useCompression) {
+    public ProtocolRequestCreator(String requestBaseUri, HashMap<String, String> headerList) {
         this.mRequestBaseUri = requestBaseUri;
         this.mHeaderList = headerList;
-        this.mUseCompression = useCompression;
     }
 
     /**
      * Creates a {@link ProtocolRequestCreator} based on forwarding info. Validates and extracts the
      * base URI for the subsequent requests.
      */
-    public static ProtocolRequestCreator create(
-            ForwardingInfo forwardingInfo, boolean useCompression) {
+    public static ProtocolRequestCreator create(ForwardingInfo forwardingInfo) {
         if (forwardingInfo.getTargetUriPrefix().isEmpty()) {
             throw new IllegalArgumentException("Missing `ForwardingInfo.target_uri_prefix`");
         }
         HashMap<String, String> extraHeaders = new HashMap<>();
         extraHeaders.putAll(forwardingInfo.getExtraRequestHeadersMap());
-        return new ProtocolRequestCreator(
-                forwardingInfo.getTargetUriPrefix(), extraHeaders, useCompression);
+        return new ProtocolRequestCreator(forwardingInfo.getTargetUriPrefix(), extraHeaders);
     }
 
     /** Creates a {@link FederatedComputeHttpRequest} with base uri and compression setting. */
@@ -70,30 +65,21 @@ public final class ProtocolRequestCreator {
     public FederatedComputeHttpRequest createProtoRequest(
             String uri,
             HttpMethod httpMethod,
-            HashMap<String, String> params,
+            HashMap<String, String> extraHeaders,
             byte[] requestBody,
             boolean isProtobufEncoded) {
         HashMap<String, String> requestHeader = new HashMap<>();
         requestHeader.putAll(mHeaderList);
+        requestHeader.putAll(extraHeaders);
 
         if (isProtobufEncoded && requestBody.length > 0) {
             requestHeader.put(CONTENT_TYPE_HDR, PROTOBUF_CONTENT_TYPE);
         }
-
-        String requestUriSuffix = uri;
-        if (!params.isEmpty()) {
-            requestUriSuffix = uri.concat("?");
-            for (String key : params.keySet()) {
-                requestUriSuffix = requestUriSuffix.concat(String.join("=", key, params.get(key)));
-            }
-        }
-
         return FederatedComputeHttpRequest.create(
-                joinBaseUriWithSuffix(mRequestBaseUri, requestUriSuffix),
+                joinBaseUriWithSuffix(mRequestBaseUri, uri),
                 httpMethod,
                 requestHeader,
-                requestBody,
-                mUseCompression);
+                requestBody);
     }
 
     private String joinBaseUriWithSuffix(String baseUri, String suffix) {
diff --git a/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java b/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
index cfd38970..eb141c65 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/scheduling/FederatedComputeJobManager.java
@@ -99,6 +99,52 @@ public class FederatedComputeJobManager {
         return sSingletonInstance;
     }
 
+    /** We enforce device idle, battery not low and unmetered network training constraints. */
+    private static byte[] buildTrainingConstraints() {
+        FlatBufferBuilder builder = new FlatBufferBuilder();
+        builder.finish(
+                TrainingConstraints.createTrainingConstraints(
+                        builder,
+                        /** requiresSchedulerIdle= */
+                        true,
+                        /** requiresSchedulerBatteryNotLow= */
+                        true,
+                        /** requiresSchedulerUnmeteredNetwork= */
+                        true));
+        return builder.sizedByteArray();
+    }
+
+    private static byte[] buildDefaultTrainingInterval() {
+        FlatBufferBuilder builder = new FlatBufferBuilder();
+        builder.finish(
+                TrainingIntervalOptions.createTrainingIntervalOptions(
+                        builder, SchedulingMode.ONE_TIME, 0));
+        return builder.sizedByteArray();
+    }
+
+    private static byte[] buildTrainingIntervalOptions(
+            @Nullable TrainingInterval trainingInterval) {
+        if (trainingInterval == null) {
+            return buildDefaultTrainingInterval();
+        }
+
+        FlatBufferBuilder builder = new FlatBufferBuilder();
+        builder.finish(
+                TrainingIntervalOptions.createTrainingIntervalOptions(
+                        builder,
+                        convertSchedulingMode(trainingInterval.getSchedulingMode()),
+                        trainingInterval.getMinimumIntervalMillis()));
+
+        return builder.sizedByteArray();
+    }
+
+    private static boolean trainingIntervalChanged(
+            TrainingOptions newTaskOptions, FederatedTrainingTask existingTask) {
+        byte[] incomingTrainingIntervalOptions =
+                buildTrainingIntervalOptions(newTaskOptions.getTrainingInterval());
+        return !Arrays.equals(incomingTrainingIntervalOptions, existingTask.intervalOptions());
+    }
+
     /**
      * Called when a client indicates via the client API that a task with the given parameters
      * should be scheduled.
@@ -327,45 +373,6 @@ public class FederatedComputeJobManager {
         return mJobSchedulerHelper.scheduleTask(mContext, newTask);
     }
 
-    /** We enforce device idle, battery not low and unmetered network training constraints. */
-    private static byte[] buildTrainingConstraints() {
-        FlatBufferBuilder builder = new FlatBufferBuilder();
-        builder.finish(
-                TrainingConstraints.createTrainingConstraints(
-                        builder,
-                        /** requiresSchedulerIdle= */
-                        true,
-                        /** requiresSchedulerBatteryNotLow= */
-                        true,
-                        /** requiresSchedulerUnmeteredNetwork= */
-                        true));
-        return builder.sizedByteArray();
-    }
-
-    private static byte[] buildDefaultTrainingInterval() {
-        FlatBufferBuilder builder = new FlatBufferBuilder();
-        builder.finish(
-                TrainingIntervalOptions.createTrainingIntervalOptions(
-                        builder, SchedulingMode.ONE_TIME, 0));
-        return builder.sizedByteArray();
-    }
-
-    private static byte[] buildTrainingIntervalOptions(
-            @Nullable TrainingInterval trainingInterval) {
-        if (trainingInterval == null) {
-            return buildDefaultTrainingInterval();
-        }
-
-        FlatBufferBuilder builder = new FlatBufferBuilder();
-        builder.finish(
-                TrainingIntervalOptions.createTrainingIntervalOptions(
-                        builder,
-                        convertSchedulingMode(trainingInterval.getSchedulingMode()),
-                        trainingInterval.getMinimumIntervalMillis()));
-
-        return builder.sizedByteArray();
-    }
-
     private boolean detectKeyParametersChanged(
             TrainingOptions newTaskOptions, FederatedTrainingTask existingTask) {
         // Check if the task previously had a different population name.
@@ -389,11 +396,4 @@ public class FederatedComputeJobManager {
         }
         return populationChanged || trainingIntervalChanged;
     }
-
-    private static boolean trainingIntervalChanged(
-            TrainingOptions newTaskOptions, FederatedTrainingTask existingTask) {
-        byte[] incomingTrainingIntervalOptions =
-                buildTrainingIntervalOptions(newTaskOptions.getTrainingInterval());
-        return !Arrays.equals(incomingTrainingIntervalOptions, existingTask.intervalOptions());
-    }
 }
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
index c991b916..4244aff4 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/FederatedComputeWorker.java
@@ -18,6 +18,8 @@ package com.android.federatedcompute.services.training;
 
 import static com.android.federatedcompute.services.common.Constants.CLIENT_ONLY_PLAN_FILE_NAME;
 import static com.android.federatedcompute.services.common.Constants.ISOLATED_TRAINING_SERVICE_NAME;
+import static com.android.federatedcompute.services.common.Constants.TRACE_WORKER_RUN_FL_COMPUTATION;
+import static com.android.federatedcompute.services.common.Constants.TRACE_WORKER_START_TRAINING_RUN;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getBackgroundExecutor;
 import static com.android.federatedcompute.services.common.FederatedComputeExecutors.getLightweightExecutor;
 import static com.android.federatedcompute.services.common.FileUtils.createTempFile;
@@ -33,6 +35,7 @@ import android.federatedcompute.common.ClientConstants;
 import android.federatedcompute.common.ExampleConsumption;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
+import android.os.Trace;
 
 import androidx.concurrent.futures.CallbackToFutureAdapter;
 
@@ -57,6 +60,7 @@ import com.android.internal.util.Preconditions;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.FluentFuture;
+import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
@@ -66,6 +70,9 @@ import com.google.intelligence.fcp.client.RetryInfo;
 import com.google.intelligence.fcp.client.engine.TaskRetry;
 import com.google.internal.federated.plan.ClientOnlyPlan;
 import com.google.internal.federated.plan.ExampleSelector;
+import com.google.internal.federatedcompute.v1.RejectionInfo;
+import com.google.internal.federatedcompute.v1.RetryWindow;
+import com.google.protobuf.Duration;
 import com.google.protobuf.InvalidProtocolBufferException;
 
 import java.io.IOException;
@@ -157,6 +164,7 @@ public class FederatedComputeWorker {
             int jobId, FederatedTrainingTask trainingTask) {
         synchronized (mLock) {
             // Only allows one concurrent job running.
+            Trace.beginAsyncSection(TRACE_WORKER_START_TRAINING_RUN, jobId);
             TrainingRun run = new TrainingRun(jobId, trainingTask);
             mActiveRun = run;
             ListenableFuture<FLRunnerResult> runCompletedFuture = doTraining(run);
@@ -165,6 +173,8 @@ public class FederatedComputeWorker {
                             .call(
                                     () -> {
                                         unBindServicesIfNecessary(run);
+                                        Trace.endAsyncSection(
+                                                TRACE_WORKER_START_TRAINING_RUN, jobId);
                                         return null;
                                     },
                                     mInjector.getBgExecutor());
@@ -220,77 +230,170 @@ public class FederatedComputeWorker {
             ListenableFuture<CheckinResult> checkinResultFuture =
                     mHttpFederatedProtocol.issueCheckin();
 
-            // 2. Bind to client app implemented ExampleStoreService based on ExampleSelector.
-            ListenableFuture<IExampleStoreIterator> iteratorFuture =
-                    FluentFuture.from(checkinResultFuture)
-                            .transform(
-                                    result -> {
-                                        // Set active run's task name.
-                                        String taskName = result.getTaskAssignment().getTaskName();
-                                        Preconditions.checkArgument(
-                                                !taskName.isEmpty(),
-                                                "Task name should not be empty");
-                                        synchronized (mLock) {
-                                            mActiveRun.mTaskName = taskName;
-                                        }
-                                        return getExampleSelector(result);
-                                    },
-                                    getLightweightExecutor())
-                            .transformAsync(
-                                    selector ->
-                                            getExampleStoreIterator(
-                                                    run,
-                                                    run.mTask.appPackageName(),
-                                                    run.mTaskName,
-                                                    selector),
-                                    mInjector.getBgExecutor());
-
-            // 3. Run federated learning or federated analytic depends on task type. Federated
-            // learning job will start a new isolated process to run TFLite training.
-            ListenableFuture<ComputationResult> computationResultFuture =
-                    Futures.whenAllSucceed(checkinResultFuture, iteratorFuture)
-                            .callAsync(
-                                    () ->
-                                            runFederatedComputation(
-                                                    Futures.getDone(checkinResultFuture),
-                                                    run,
-                                                    Futures.getDone(iteratorFuture)),
-                                    mInjector.getBgExecutor());
-
-            // 4. Report computation result to federated compute server.
-            ListenableFuture<Void> reportToServerFuture =
-                    FluentFuture.from(computationResultFuture)
-                            .transformAsync(
-                                    result -> mHttpFederatedProtocol.reportResult(result),
-                                    getLightweightExecutor());
-            return FluentFuture.from(
-                    Futures.whenAllSucceed(reportToServerFuture, computationResultFuture)
-                            .call(
-                                    () -> {
-                                        ComputationResult result =
-                                                Futures.getDone(computationResultFuture);
-                                        var reportToServer = Futures.getDone(reportToServerFuture);
-                                        // 5. Publish computation result and consumed
-                                        // examples to client implemented
-                                        // ResultHandlingService.
-                                        var unused =
-                                                mResultCallbackHelper.callHandleResult(
-                                                        run.mTaskName, run.mTask, result);
-                                        return result.getFlRunnerResult();
-                                    },
-                                    mInjector.getBgExecutor()));
+            return FluentFuture.from(checkinResultFuture)
+                    .transformAsync(
+                            checkinResult -> processCheckinAndDoFlTraining(run, checkinResult),
+                            mInjector.getBgExecutor());
         } catch (Exception e) {
             return Futures.immediateFailedFuture(e);
         }
     }
 
+    @androidx.annotation.NonNull
+    private ListenableFuture<FLRunnerResult> processCheckinAndDoFlTraining(
+            TrainingRun run, CheckinResult checkinResult) {
+        // Stop processing if have rejection Info
+        if (checkinResult.getRejectionInfo() != null) {
+            mJobManager.onTrainingCompleted(
+                    run.mTask.jobId(),
+                    run.mTask.populationName(),
+                    run.mTask.getTrainingIntervalOptions(),
+                    buildTaskRetry(checkinResult.getRejectionInfo()),
+                    ContributionResult.FAIL);
+            return Futures.immediateFuture(null);
+        }
+        // 2. Bind to client app implemented ExampleStoreService based on ExampleSelector.
+        // Set active run's task name.
+        String taskName = checkinResult.getTaskAssignment().getTaskName();
+        Preconditions.checkArgument(!taskName.isEmpty(), "Task name should not be empty");
+        synchronized (mLock) {
+            mActiveRun.mTaskName = taskName;
+        }
+        ListenableFuture<IExampleStoreIterator> iteratorFuture =
+                getExampleStoreIterator(
+                        run,
+                        run.mTask.appPackageName(),
+                        run.mTaskName,
+                        getExampleSelector(checkinResult));
+        // report failure to server if getting iterator failed with any exception.
+        FutureCallback<Object> serverFailureReportCallback = getServerFailureReportCallback();
+        Futures.addCallback(
+                iteratorFuture, serverFailureReportCallback, getLightweightExecutor());
+
+        // 3. Run federated learning or federated analytic depends on task type. Federated
+        // learning job will start a new isolated process to run TFLite training.
+        FluentFuture<ComputationResult> computationResultFuture =
+                FluentFuture.from(iteratorFuture)
+                        .transformAsync(
+                                iterator -> runFederatedComputation(checkinResult, run, iterator),
+                                mInjector.getBgExecutor());
+        // report failure to server if computation failed with any exception.
+        computationResultFuture.addCallback(
+                serverFailureReportCallback, getLightweightExecutor());
+
+        // 4. Report computation result to federated compute server.
+        ListenableFuture<RejectionInfo> reportToServerFuture =
+                computationResultFuture.transformAsync(
+                        result -> mHttpFederatedProtocol.reportResult(result),
+                        getLightweightExecutor());
+        return Futures.whenAllSucceed(reportToServerFuture, computationResultFuture)
+                .call(
+                        () -> {
+                            ComputationResult computationResult =
+                                    Futures.getDone(computationResultFuture);
+                            RejectionInfo reportToServer = Futures.getDone(reportToServerFuture);
+                            // report to Server will hold null in case of success, or rejection info
+                            // in case server answered with rejection
+                            if (reportToServer != null) {
+                                ComputationResult failedReportComputationResult =
+                                        new ComputationResult(
+                                                null,
+                                                FLRunnerResult.newBuilder()
+                                                        .mergeFrom(
+                                                                computationResult
+                                                                        .getFlRunnerResult())
+                                                        .setContributionResult(
+                                                                ContributionResult.FAIL)
+                                                        .build(),
+                                                null);
+                                var unused =
+                                        mResultCallbackHelper.callHandleResult(
+                                                run.mTaskName,
+                                                run.mTask,
+                                                failedReportComputationResult);
+                                mJobManager.onTrainingCompleted(
+                                        run.mTask.jobId(),
+                                        run.mTask.populationName(),
+                                        run.mTask.getTrainingIntervalOptions(),
+                                        buildTaskRetry(reportToServer),
+                                        ContributionResult.FAIL);
+                                return null;
+                            }
+                            // 5. Publish computation result and consumed
+                            // examples to client implemented
+                            // ResultHandlingService.
+                            var unused =
+                                    mResultCallbackHelper.callHandleResult(
+                                            run.mTaskName, run.mTask, computationResult);
+                            return computationResult.getFlRunnerResult();
+                        },
+                        mInjector.getBgExecutor());
+    }
+
+    @androidx.annotation.NonNull
+    private FutureCallback<Object> getServerFailureReportCallback() {
+        return new FutureCallback<Object>() {
+            volatile int mNumberOfInvocations = 0;
+
+            @Override
+            public void onSuccess(Object unused) {
+                // do nothing.
+            }
+
+            // We do not want race condition and repeating reporting failures from computation
+            // failed future in right before case Example Store iterator failed.
+            // Thus method is synchronised.
+            @Override
+            public synchronized void onFailure(Throwable throwable) {
+                if (mNumberOfInvocations < 1) {
+                    LogUtil.d(
+                            TAG,
+                            "Training failed. Reporting failure result to server due to exception.",
+                            throwable);
+                    ComputationResult failedReportComputationResult =
+                            new ComputationResult(
+                                    null,
+                                    FLRunnerResult.newBuilder()
+                                            .setContributionResult(ContributionResult.FAIL)
+                                            .setErrorMessage(throwable.getMessage())
+                                            .build(),
+                                    null);
+                    var unused = mHttpFederatedProtocol.reportResult(failedReportComputationResult);
+                }
+                mNumberOfInvocations++;
+            }
+        };
+    }
+
+    private static TaskRetry buildTaskRetry(RejectionInfo rejectionInfo) {
+        TaskRetry.Builder taskRetryBuilder =
+                TaskRetry.newBuilder();
+        if (rejectionInfo.hasRetryWindow()) {
+            RetryWindow retryWindow =
+                    rejectionInfo.getRetryWindow();
+            Duration delayMin = retryWindow.getDelayMin();
+            // convert rejection info seconds and nanoseconds to
+            // retry milliseconds
+            taskRetryBuilder.setDelayMin(
+                    delayMin.getSeconds() * 1000
+                            + delayMin.getNanos() / 1000000);
+            Duration delayMax = retryWindow.getDelayMax();
+            taskRetryBuilder.setDelayMax(
+                    delayMax.getSeconds() * 1000
+                            + delayMax.getNanos() / 1000000);
+        }
+        return taskRetryBuilder.build();
+    }
+
     /**
      * Completes the running job , schedule recurrent job, and unbind from ExampleStoreService and
      * ResultHandlingService etc.
      */
     public void finish(FLRunnerResult flRunnerResult) {
         TaskRetry taskRetry = null;
+        ContributionResult contributionResult = ContributionResult.UNSPECIFIED;
         if (flRunnerResult != null) {
+            contributionResult = flRunnerResult.getContributionResult();
             if (flRunnerResult.hasRetryInfo()) {
                 RetryInfo retryInfo = flRunnerResult.getRetryInfo();
                 long delay = retryInfo.getMinimumDelay().getSeconds() * 1000L;
@@ -303,7 +406,7 @@ public class FederatedComputeWorker {
                 LogUtil.i(TAG, "Finished with task retry= %s", taskRetry);
             }
         }
-        finish(taskRetry, flRunnerResult.getContributionResult(), true);
+        finish(taskRetry, contributionResult, true);
     }
 
     /**
@@ -393,6 +496,7 @@ public class FederatedComputeWorker {
             CheckinResult checkinResult,
             String outputCheckpointFile,
             IExampleStoreIterator iterator) {
+        Trace.beginAsyncSection(TRACE_WORKER_RUN_FL_COMPUTATION, 0);
         ParcelFileDescriptor outputCheckpointFd =
                 createTempFileDescriptor(
                         outputCheckpointFile, ParcelFileDescriptor.MODE_READ_WRITE);
@@ -458,6 +562,7 @@ public class FederatedComputeWorker {
                                     unbindFromIsolatedTrainingService();
                                     run.mIsolatedTrainingService = null;
                                 }
+                                Trace.endAsyncSection(TRACE_WORKER_RUN_FL_COMPUTATION, 0);
                                 return computationResult;
                             },
                             getLightweightExecutor());
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
index ebaf7ef0..79966d99 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/IsolatedTrainingServiceImpl.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.training;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING;
+
 import android.annotation.NonNull;
 import android.federatedcompute.aidl.IExampleStoreIterator;
 import android.federatedcompute.common.ClientConstants;
@@ -23,6 +25,7 @@ import android.federatedcompute.common.ExampleConsumption;
 import android.os.Bundle;
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
+import android.os.Trace;
 
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.common.Constants;
@@ -70,6 +73,7 @@ public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
     public void runFlTraining(@NonNull Bundle params, @NonNull ITrainingResultCallback callback) {
         Objects.requireNonNull(params);
         Objects.requireNonNull(callback);
+        Trace.beginAsyncSection(TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING, 0);
 
         IExampleStoreIterator exampleStoreIteratorBinder =
                 IExampleStoreIterator.Stub.asInterface(
@@ -135,16 +139,21 @@ public class IsolatedTrainingServiceImpl extends IIsolatedTrainingService.Stub {
                         bundle.putByteArray(Constants.EXTRA_FL_RUNNER_RESULT, result.toByteArray());
                         ArrayList<ExampleConsumption> exampleConsumptionArrayList =
                                 recorder.finishRecordingAndGet();
+                        int numExamples = 0;
+                        for (ExampleConsumption exampleConsumption : exampleConsumptionArrayList) {
+                            numExamples += exampleConsumption.getExampleCount();
+                        }
                         LogUtil.i(
                                 TAG,
                                 "training task %s: result %s, used %d examples",
                                 populationName,
                                 result.toString(),
-                                exampleConsumptionArrayList.size());
+                                numExamples);
                         bundle.putParcelableArrayList(
                                 ClientConstants.EXTRA_EXAMPLE_CONSUMPTION_LIST,
                                 exampleConsumptionArrayList);
                         sendResult(bundle, callback);
+                        Trace.endAsyncSection(TRACE_ISOLATED_PROCESS_RUN_FL_TRAINING, 0);
                     }
 
                     @Override
diff --git a/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java b/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
index 33ec656f..30fa4ce1 100644
--- a/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
+++ b/federatedcompute/src/com/android/federatedcompute/services/training/jni/FlRunnerWrapper.java
@@ -16,6 +16,10 @@
 
 package com.android.federatedcompute.services.training.jni;
 
+import static com.android.federatedcompute.services.common.Constants.TRACE_NATIVE_RUN_FEDERATED_COMPUTATION;
+
+import android.os.Trace;
+
 import com.android.federatedcompute.internal.util.LogUtil;
 import com.android.federatedcompute.services.examplestore.ExampleIterator;
 import com.android.federatedcompute.services.training.util.ListenableSupplier;
@@ -57,6 +61,7 @@ public final class FlRunnerWrapper implements Closeable {
             ClientOnlyPlan clientOnlyPlan,
             String checkpointInputFileName,
             String checkpointOutputFileName) {
+        Trace.beginAsyncSection(TRACE_NATIVE_RUN_FEDERATED_COMPUTATION, 0);
         SimpleTaskEnvironmentImpl simpleTaskEnv =
                 new SimpleTaskEnvironmentImpl(mInterruptionFlag, mExampleIterator);
         byte[] flRunnerResultSerialized =
@@ -77,6 +82,8 @@ public final class FlRunnerWrapper implements Closeable {
             // recover from it.
             LogUtil.e(TAG, "Cannot parse FLRunnerResult", e);
             throw new IllegalArgumentException(e);
+        } finally {
+            Trace.endAsyncSection(TRACE_NATIVE_RUN_FEDERATED_COMPUTATION, 0);
         }
     }
 
diff --git a/framework/api/current.txt b/framework/api/current.txt
index e07681c6..87e868ae 100644
--- a/framework/api/current.txt
+++ b/framework/api/current.txt
@@ -12,11 +12,8 @@ package android.adservices.ondevicepersonalization {
     method @NonNull public java.util.Map<java.lang.String,byte[]> getData();
   }
 
-  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class DownloadCompletedOutput implements android.os.Parcelable {
-    method public int describeContents();
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class DownloadCompletedOutput {
     method @NonNull public java.util.List<java.lang.String> getRetainedKeys();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.DownloadCompletedOutput> CREATOR;
   }
 
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public static final class DownloadCompletedOutput.Builder {
@@ -26,10 +23,17 @@ package android.adservices.ondevicepersonalization {
     method @NonNull public android.adservices.ondevicepersonalization.DownloadCompletedOutput.Builder setRetainedKeys(@NonNull java.util.List<java.lang.String>);
   }
 
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class EventInput {
+    method @NonNull public android.os.PersistableBundle getParameters();
+    method @Nullable public android.adservices.ondevicepersonalization.RequestLogRecord getRequestLogRecord();
+  }
+
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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;
@@ -39,29 +43,34 @@ package android.adservices.ondevicepersonalization {
     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(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class EventOutput {
+    method @Nullable public android.adservices.ondevicepersonalization.EventLogRecord getEventLogRecord();
+  }
+
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class ExecuteInput implements android.os.Parcelable {
-    method public int describeContents();
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class ExecuteInput {
     method @NonNull public String getAppPackageName();
     method @NonNull public android.os.PersistableBundle getAppParams();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.ExecuteInput> CREATOR;
   }
 
-  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class ExecuteOutput implements android.os.Parcelable {
-    method public int describeContents();
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.ExecuteOutput> CREATOR;
   }
 
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public static final class ExecuteOutput.Builder {
@@ -72,10 +81,32 @@ package android.adservices.ondevicepersonalization {
     method @NonNull public android.adservices.ondevicepersonalization.ExecuteOutput.Builder setRequestLogRecord(@NonNull android.adservices.ondevicepersonalization.RequestLogRecord);
   }
 
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class FederatedComputeInput {
+    method @NonNull public String getPopulationName();
+  }
+
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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);
@@ -84,8 +115,10 @@ package android.adservices.ondevicepersonalization {
 
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public interface KeyValueStore {
@@ -93,6 +126,11 @@ package android.adservices.ondevicepersonalization {
     method @NonNull @WorkerThread public java.util.Set<java.lang.String> keySet();
   }
 
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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);
@@ -104,26 +142,21 @@ package android.adservices.ondevicepersonalization {
   }
 
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class RenderInput implements android.os.Parcelable {
-    method public int describeContents();
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.RenderInput> CREATOR;
   }
 
-  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class RenderOutput implements android.os.Parcelable {
-    method public int describeContents();
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class RenderOutput {
     method @Nullable public String getContent();
     method @Nullable public String getTemplateId();
     method @NonNull public android.os.PersistableBundle getTemplateParams();
-    method public void writeToParcel(@NonNull android.os.Parcel, int);
-    field @NonNull public static final android.os.Parcelable.Creator<android.adservices.ondevicepersonalization.RenderOutput> CREATOR;
   }
 
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public static final class RenderOutput.Builder {
@@ -151,6 +184,7 @@ package android.adservices.ondevicepersonalization {
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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;
   }
@@ -168,6 +202,40 @@ package android.adservices.ondevicepersonalization {
   @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public class SurfacePackageToken {
   }
 
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class TrainingExamplesInput {
+    method @NonNull public String getPopulationName();
+    method @Nullable public byte[] getResumptionToken();
+    method @NonNull public String getTaskName();
+  }
+
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS) public final class TrainingExamplesOutput {
+    method @NonNull public java.util.List<byte[]> getResumptionTokens();
+    method @NonNull public java.util.List<byte[]> getTrainingExamples();
+  }
+
+  @FlaggedApi(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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(android.adservices.ondevicepersonalization.Constants.KEY_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();
diff --git a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.aidl b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.aidl
deleted file mode 100644
index 8ebe46ef..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.aidl
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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 android.adservices.ondevicepersonalization;
-
-parcelable DownloadCompletedOutput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
index 10d953d0..5c503fcb 100644
--- a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java
@@ -20,7 +20,6 @@ import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ON
 
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
-import android.os.Parcelable;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
@@ -35,7 +34,7 @@ import java.util.List;
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
-public final class DownloadCompletedOutput implements Parcelable {
+public final class DownloadCompletedOutput {
     /**
      * The keys to be retained in the REMOTE_DATA table. Any existing keys that are not
      * present in this list are removed from the table.
@@ -104,50 +103,6 @@ public final class DownloadCompletedOutput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        dest.writeStringList(mRetainedKeys);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ DownloadCompletedOutput(@NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        List<String> retainedKeys = new java.util.ArrayList<>();
-        in.readStringList(retainedKeys);
-
-        this.mRetainedKeys = retainedKeys;
-        AnnotationValidations.validate(
-                NonNull.class, null, mRetainedKeys);
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<DownloadCompletedOutput> CREATOR
-            = new Parcelable.Creator<DownloadCompletedOutput>() {
-        @Override
-        public DownloadCompletedOutput[] newArray(int size) {
-            return new DownloadCompletedOutput[size];
-        }
-
-        @Override
-        public DownloadCompletedOutput createFromParcel(@NonNull android.os.Parcel in) {
-            return new DownloadCompletedOutput(in);
-        }
-    };
-
     /**
      * A builder for {@link DownloadCompletedOutput}
      */
@@ -205,10 +160,10 @@ public final class DownloadCompletedOutput implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1696972554365L,
+            time = 1698862918590L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutput.java",
-            inputSignatures = "private @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"retainedKey\") @android.annotation.NonNull java.util.List<java.lang.String> mRetainedKeys\nclass DownloadCompletedOutput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = "private @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"retainedKey\") @android.annotation.NonNull java.util.List<java.lang.String> mRetainedKeys\nclass DownloadCompletedOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutputParcel.java b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutputParcel.java
new file mode 100644
index 00000000..7a906f1d
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutputParcel.java
@@ -0,0 +1,142 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.os.Parcelable;
+
+import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Parcelable version of {@link DownloadCompletedOutput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genBuilder =  false)
+public final class DownloadCompletedOutputParcel implements Parcelable {
+    /**
+     * The keys to be retained in the REMOTE_DATA table. Any existing keys that are not
+     * present in this list are removed from the table.
+     */
+    @NonNull private List<String> mRetainedKeys = Collections.emptyList();
+
+    /** @hide */
+    public DownloadCompletedOutputParcel(@NonNull DownloadCompletedOutput value) {
+        this(value.getRetainedKeys());
+    }
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    /**
+     * Creates a new DownloadCompletedOutputParcel.
+     *
+     * @param retainedKeys
+     *   The keys to be retained in the REMOTE_DATA table. Any existing keys that are not
+     *   present in this list are removed from the table.
+     */
+    @DataClass.Generated.Member
+    public DownloadCompletedOutputParcel(
+            @NonNull List<String> retainedKeys) {
+        this.mRetainedKeys = retainedKeys;
+        AnnotationValidations.validate(
+                NonNull.class, null, mRetainedKeys);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The keys to be retained in the REMOTE_DATA table. Any existing keys that are not
+     * present in this list are removed from the table.
+     */
+    @DataClass.Generated.Member
+    public @NonNull List<String> getRetainedKeys() {
+        return mRetainedKeys;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        dest.writeStringList(mRetainedKeys);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ DownloadCompletedOutputParcel(@NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        List<String> retainedKeys = new java.util.ArrayList<>();
+        in.readStringList(retainedKeys);
+
+        this.mRetainedKeys = retainedKeys;
+        AnnotationValidations.validate(
+                NonNull.class, null, mRetainedKeys);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @NonNull Parcelable.Creator<DownloadCompletedOutputParcel> CREATOR
+            = new Parcelable.Creator<DownloadCompletedOutputParcel>() {
+        @Override
+        public DownloadCompletedOutputParcel[] newArray(int size) {
+            return new DownloadCompletedOutputParcel[size];
+        }
+
+        @Override
+        public DownloadCompletedOutputParcel createFromParcel(@NonNull android.os.Parcel in) {
+            return new DownloadCompletedOutputParcel(in);
+        }
+    };
+
+    @DataClass.Generated(
+            time = 1698783477713L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/DownloadCompletedOutputParcel.java",
+            inputSignatures = "private @android.annotation.NonNull java.util.List<java.lang.String> mRetainedKeys\nclass DownloadCompletedOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genBuilder=false)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventInput.aidl b/framework/java/android/adservices/ondevicepersonalization/EventInput.aidl
deleted file mode 100644
index 156f47e7..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/EventInput.aidl
+++ /dev/null
@@ -1,19 +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.adservices.ondevicepersonalization;
-
-parcelable EventInput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventInput.java b/framework/java/android/adservices/ondevicepersonalization/EventInput.java
index 8540aaf3..33057ad6 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventInput.java
@@ -16,9 +16,11 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Parcelable;
 import android.os.PersistableBundle;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
@@ -27,10 +29,10 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 /**
  * The input data for {@link
  * IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}.
- * @hide
  */
-@DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
-public final class EventInput implements Parcelable {
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
+@DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
+public final class EventInput {
     /**
      * The {@link RequestLogRecord} that was returned as a result of
      * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
@@ -44,6 +46,11 @@ public final class EventInput implements Parcelable {
      */
     @NonNull private PersistableBundle mParameters = PersistableBundle.EMPTY;
 
+    /** @hide */
+    public EventInput(@NonNull EventInputParcel parcel) {
+        this(parcel.getRequestLogRecord(), parcel.getParameters());
+    }
+
 
 
     // Code below generated by codegen v1.0.23.
@@ -59,8 +66,20 @@ public final class EventInput implements Parcelable {
     //@formatter:off
 
 
+    /**
+     * Creates a new EventInput.
+     *
+     * @param requestLogRecord
+     *   The {@link RequestLogRecord} that was returned as a result of
+     *   {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     * @param parameters
+     *   The Event URL parameters that the service passed to {@link
+     *   EventUrlProvider#createEventTrackingUrlWithResponse(PersistableBundle, byte[], String)}
+     *   or {@link EventUrlProvider#createEventTrackingUrlWithRedirect(PersistableBundle, Uri)}.
+     * @hide
+     */
     @DataClass.Generated.Member
-    /* package-private */ EventInput(
+    public EventInput(
             @Nullable RequestLogRecord requestLogRecord,
             @NonNull PersistableBundle parameters) {
         this.mRequestLogRecord = requestLogRecord;
@@ -119,127 +138,11 @@ public final class EventInput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        byte flg = 0;
-        if (mRequestLogRecord != null) flg |= 0x1;
-        dest.writeByte(flg);
-        if (mRequestLogRecord != null) dest.writeTypedObject(mRequestLogRecord, flags);
-        dest.writeTypedObject(mParameters, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ EventInput(@NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        byte flg = in.readByte();
-        RequestLogRecord requestLogRecord = (flg & 0x1) == 0 ? null : (RequestLogRecord) in.readTypedObject(RequestLogRecord.CREATOR);
-        PersistableBundle parameters = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
-
-        this.mRequestLogRecord = requestLogRecord;
-        this.mParameters = parameters;
-        AnnotationValidations.validate(
-                NonNull.class, null, mParameters);
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<EventInput> CREATOR
-            = new Parcelable.Creator<EventInput>() {
-        @Override
-        public EventInput[] newArray(int size) {
-            return new EventInput[size];
-        }
-
-        @Override
-        public EventInput createFromParcel(@NonNull android.os.Parcel in) {
-            return new EventInput(in);
-        }
-    };
-
-    /**
-     * A builder for {@link EventInput}
-     * @hide
-     */
-    @SuppressWarnings("WeakerAccess")
-    @DataClass.Generated.Member
-    public static final class Builder {
-
-        private @Nullable RequestLogRecord mRequestLogRecord;
-        private @NonNull PersistableBundle mParameters;
-
-        private long mBuilderFieldsSet = 0L;
-
-        public Builder() {
-        }
-
-        /**
-         * The {@link RequestLogRecord} that was returned as a result of
-         * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
-         */
-        @DataClass.Generated.Member
-        public @NonNull Builder setRequestLogRecord(@NonNull RequestLogRecord value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x1;
-            mRequestLogRecord = value;
-            return this;
-        }
-
-        /**
-         * The Event URL parameters that the service passed to {@link
-         * EventUrlProvider#createEventTrackingUrlWithResponse(PersistableBundle, byte[], String)}
-         * or {@link EventUrlProvider#createEventTrackingUrlWithRedirect(PersistableBundle, Uri)}.
-         */
-        @DataClass.Generated.Member
-        public @NonNull Builder setParameters(@NonNull PersistableBundle value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x2;
-            mParameters = value;
-            return this;
-        }
-
-        /** Builds the instance. This builder should not be touched after calling this! */
-        public @NonNull EventInput build() {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x4; // Mark builder used
-
-            if ((mBuilderFieldsSet & 0x1) == 0) {
-                mRequestLogRecord = null;
-            }
-            if ((mBuilderFieldsSet & 0x2) == 0) {
-                mParameters = PersistableBundle.EMPTY;
-            }
-            EventInput o = new EventInput(
-                    mRequestLogRecord,
-                    mParameters);
-            return o;
-        }
-
-        private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x4) != 0) {
-                throw new IllegalStateException(
-                        "This Builder should not be reused. Use a new Builder instance instead");
-            }
-        }
-    }
-
     @DataClass.Generated(
-            time = 1697062257720L,
+            time = 1698882321696L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventInput.java",
-            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @android.annotation.NonNull android.os.PersistableBundle mParameters\nclass EventInput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genHiddenBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @android.annotation.NonNull android.os.PersistableBundle mParameters\nclass EventInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=false, genHiddenConstructor=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventInputParcel.java b/framework/java/android/adservices/ondevicepersonalization/EventInputParcel.java
new file mode 100644
index 00000000..7434064d
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/EventInputParcel.java
@@ -0,0 +1,220 @@
+/*
+ * 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.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+
+import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+/**
+ * Parcelable version of {@link EventInput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genHiddenBuilder = true)
+public final class EventInputParcel implements Parcelable {
+    /**
+     * The {@link RequestLogRecord} that was returned as a result of
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     */
+    @Nullable private RequestLogRecord mRequestLogRecord = null;
+
+    /**
+     * The Event URL parameters that the service passed to {@link
+     * EventUrlProvider#createEventTrackingUrlWithResponse(PersistableBundle, byte[], String)}
+     * or {@link EventUrlProvider#createEventTrackingUrlWithRedirect(PersistableBundle, Uri)}.
+     */
+    @NonNull private PersistableBundle mParameters = PersistableBundle.EMPTY;
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventInputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    @DataClass.Generated.Member
+    /* package-private */ EventInputParcel(
+            @Nullable RequestLogRecord requestLogRecord,
+            @NonNull PersistableBundle parameters) {
+        this.mRequestLogRecord = requestLogRecord;
+        this.mParameters = parameters;
+        AnnotationValidations.validate(
+                NonNull.class, null, mParameters);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The {@link RequestLogRecord} that was returned as a result of
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     */
+    @DataClass.Generated.Member
+    public @Nullable RequestLogRecord getRequestLogRecord() {
+        return mRequestLogRecord;
+    }
+
+    /**
+     * The Event URL parameters that the service passed to {@link
+     * EventUrlProvider#createEventTrackingUrlWithResponse(PersistableBundle, byte[], String)}
+     * or {@link EventUrlProvider#createEventTrackingUrlWithRedirect(PersistableBundle, Uri)}.
+     */
+    @DataClass.Generated.Member
+    public @NonNull PersistableBundle getParameters() {
+        return mParameters;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        byte flg = 0;
+        if (mRequestLogRecord != null) flg |= 0x1;
+        dest.writeByte(flg);
+        if (mRequestLogRecord != null) dest.writeTypedObject(mRequestLogRecord, flags);
+        dest.writeTypedObject(mParameters, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ EventInputParcel(@NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        byte flg = in.readByte();
+        RequestLogRecord requestLogRecord = (flg & 0x1) == 0 ? null : (RequestLogRecord) in.readTypedObject(RequestLogRecord.CREATOR);
+        PersistableBundle parameters = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
+
+        this.mRequestLogRecord = requestLogRecord;
+        this.mParameters = parameters;
+        AnnotationValidations.validate(
+                NonNull.class, null, mParameters);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @NonNull Parcelable.Creator<EventInputParcel> CREATOR
+            = new Parcelable.Creator<EventInputParcel>() {
+        @Override
+        public EventInputParcel[] newArray(int size) {
+            return new EventInputParcel[size];
+        }
+
+        @Override
+        public EventInputParcel createFromParcel(@NonNull android.os.Parcel in) {
+            return new EventInputParcel(in);
+        }
+    };
+
+    /**
+     * A builder for {@link EventInputParcel}
+     * @hide
+     */
+    @SuppressWarnings("WeakerAccess")
+    @DataClass.Generated.Member
+    public static final class Builder {
+
+        private @Nullable RequestLogRecord mRequestLogRecord;
+        private @NonNull PersistableBundle mParameters;
+
+        private long mBuilderFieldsSet = 0L;
+
+        public Builder() {
+        }
+
+        /**
+         * The {@link RequestLogRecord} that was returned as a result of
+         * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setRequestLogRecord(@NonNull RequestLogRecord value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x1;
+            mRequestLogRecord = value;
+            return this;
+        }
+
+        /**
+         * The Event URL parameters that the service passed to {@link
+         * EventUrlProvider#createEventTrackingUrlWithResponse(PersistableBundle, byte[], String)}
+         * or {@link EventUrlProvider#createEventTrackingUrlWithRedirect(PersistableBundle, Uri)}.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setParameters(@NonNull PersistableBundle value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x2;
+            mParameters = value;
+            return this;
+        }
+
+        /** Builds the instance. This builder should not be touched after calling this! */
+        public @NonNull EventInputParcel build() {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x4; // Mark builder used
+
+            if ((mBuilderFieldsSet & 0x1) == 0) {
+                mRequestLogRecord = null;
+            }
+            if ((mBuilderFieldsSet & 0x2) == 0) {
+                mParameters = PersistableBundle.EMPTY;
+            }
+            EventInputParcel o = new EventInputParcel(
+                    mRequestLogRecord,
+                    mParameters);
+            return o;
+        }
+
+        private void checkNotUsed() {
+            if ((mBuilderFieldsSet & 0x4) != 0) {
+                throw new IllegalStateException(
+                        "This Builder should not be reused. Use a new Builder instance instead");
+            }
+        }
+    }
+
+    @DataClass.Generated(
+            time = 1698875208124L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventInputParcel.java",
+            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @android.annotation.NonNull android.os.PersistableBundle mParameters\nclass EventInputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genHiddenBuilder=true)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java b/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
index a65b2aa4..bb1c6101 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java
@@ -28,6 +28,8 @@ import android.os.Parcelable;
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
+import java.time.Instant;
+
 // TODO(b/289102463): Add a link to the public doc for the EVENTS table when available.
 /**
  * Data to be logged in the EVENTS table.
@@ -75,15 +77,20 @@ public final class EventLogRecord implements Parcelable {
      * The existing {@link RequestLogRecord} that this payload should be associated with. In an
      * implementation of
      * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}, this should be
-     * set to a value returned by {@link LogReader#getRequests(long, long)}. In an implementation
-     * of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}, this should be
-     * set to {@code null} because the payload will be automatically associated with the current
-     * {@link RequestLogRecord}.
+     * set to a value returned by {@link LogReader#getRequests(Instant, Instant)}. In an
+     * implementation of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)},
+     * this should be set to {@code null} because the payload will be automatically associated with
+     * the current {@link RequestLogRecord}.
      *
-     * @hide
      */
     @Nullable RequestLogRecord mRequestLogRecord = null;
 
+    /**
+     * Returns the timestamp of this record.
+     */
+    @NonNull public Instant getTime() {
+        return Instant.ofEpochMilli(getTimeMillis());
+    }
 
     abstract static class BaseBuilder {
         /**
@@ -174,12 +181,10 @@ public final class EventLogRecord implements Parcelable {
      * The existing {@link RequestLogRecord} that this payload should be associated with. In an
      * implementation of
      * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}, this should be
-     * set to a value returned by {@link LogReader#getRequests(long, long)}. In an implementation
-     * of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}, this should be
-     * set to {@code null} because the payload will be automatically associated with the current
-     * {@link RequestLogRecord}.
-     *
-     * @hide
+     * set to a value returned by {@link LogReader#getRequests(Instant, Instant)}. In an
+     * implementation of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)},
+     * this should be set to {@code null} because the payload will be automatically associated with
+     * the current {@link RequestLogRecord}.
      */
     @DataClass.Generated.Member
     public @Nullable RequestLogRecord getRequestLogRecord() {
@@ -362,12 +367,10 @@ public final class EventLogRecord implements Parcelable {
          * The existing {@link RequestLogRecord} that this payload should be associated with. In an
          * implementation of
          * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}, this should be
-         * set to a value returned by {@link LogReader#getRequests(long, long)}. In an implementation
-         * of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)}, this should be
-         * set to {@code null} because the payload will be automatically associated with the current
-         * {@link RequestLogRecord}.
-         *
-         * @hide
+         * set to a value returned by {@link LogReader#getRequests(Instant, Instant)}. In an
+         * implementation of {@link IsolatedWorker#onEvent(EventInput, java.util.function.Consumer)},
+         * this should be set to {@code null} because the payload will be automatically associated with
+         * the current {@link RequestLogRecord}.
          */
         @DataClass.Generated.Member
         public @NonNull Builder setRequestLogRecord(@NonNull RequestLogRecord value) {
@@ -415,10 +418,10 @@ public final class EventLogRecord implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1697576750150L,
+            time = 1698963996102L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventLogRecord.java",
-            inputSignatures = "private @android.annotation.IntRange int mRowIndex\nprivate @android.annotation.IntRange int mType\nprivate  long mTimeMillis\n @android.annotation.Nullable android.content.ContentValues mData\n @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nclass EventLogRecord extends java.lang.Object implements [android.os.Parcelable]\npublic abstract  android.adservices.ondevicepersonalization.EventLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)\npublic abstract  android.adservices.ondevicepersonalization.EventLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []")
+            inputSignatures = "private @android.annotation.IntRange int mRowIndex\nprivate @android.annotation.IntRange int mType\nprivate  long mTimeMillis\n @android.annotation.Nullable android.content.ContentValues mData\n @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\npublic @android.annotation.NonNull java.time.Instant getTime()\nclass EventLogRecord extends java.lang.Object implements [android.os.Parcelable]\npublic abstract  android.adservices.ondevicepersonalization.EventLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)\npublic abstract  android.adservices.ondevicepersonalization.EventLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventOutput.aidl b/framework/java/android/adservices/ondevicepersonalization/EventOutput.aidl
deleted file mode 100644
index e7e0f952..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/EventOutput.aidl
+++ /dev/null
@@ -1,19 +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.adservices.ondevicepersonalization;
-
-parcelable EventOutput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventOutput.java b/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
index 82b65750..9fe978d6 100644
--- a/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/EventOutput.java
@@ -16,17 +16,19 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
 import android.annotation.Nullable;
-import android.os.Parcelable;
 
 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)
-public final class EventOutput implements Parcelable {
+public final class EventOutput {
     /**
      * An {@link EventLogRecord} to be written to the EVENTS table, if not null. Each
      * {@link EventLogRecord} is associated with a row in an existing {@link RequestLogRecord} that
@@ -94,54 +96,10 @@ public final class EventOutput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@android.annotation.NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        byte flg = 0;
-        if (mEventLogRecord != null) flg |= 0x1;
-        dest.writeByte(flg);
-        if (mEventLogRecord != null) dest.writeTypedObject(mEventLogRecord, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ EventOutput(@android.annotation.NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        byte flg = in.readByte();
-        EventLogRecord eventLogRecord = (flg & 0x1) == 0 ? null : (EventLogRecord) in.readTypedObject(EventLogRecord.CREATOR);
-
-        this.mEventLogRecord = eventLogRecord;
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @android.annotation.NonNull Parcelable.Creator<EventOutput> CREATOR
-            = new Parcelable.Creator<EventOutput>() {
-        @Override
-        public EventOutput[] newArray(int size) {
-            return new EventOutput[size];
-        }
-
-        @Override
-        public EventOutput createFromParcel(@android.annotation.NonNull android.os.Parcel in) {
-            return new EventOutput(in);
-        }
-    };
-
     /**
      * A builder for {@link EventOutput}
      */
+    @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
     @SuppressWarnings("WeakerAccess")
     @DataClass.Generated.Member
     public static final class Builder {
@@ -188,10 +146,10 @@ public final class EventOutput implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1696369232183L,
+            time = 1698882423016L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventOutput.java",
-            inputSignatures = " @android.annotation.Nullable android.adservices.ondevicepersonalization.EventLogRecord mEventLogRecord\nclass EventOutput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = " @android.annotation.Nullable android.adservices.ondevicepersonalization.EventLogRecord mEventLogRecord\nclass EventOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/EventOutputParcel.java b/framework/java/android/adservices/ondevicepersonalization/EventOutputParcel.java
new file mode 100644
index 00000000..5432a6c4
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/EventOutputParcel.java
@@ -0,0 +1,141 @@
+/*
+ * 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.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Parcelable;
+
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+/**
+ * Parcelable version of {@link EventOutput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genBuilder = false)
+public final class EventOutputParcel implements Parcelable {
+    /**
+     * An {@link EventLogRecord} to be written to the EVENTS table, if not null. Each
+     * {@link EventLogRecord} is associated with a row in an existing {@link RequestLogRecord} that
+     * has been written to the REQUESTS table.
+     */
+    @Nullable EventLogRecord mEventLogRecord = null;
+
+    /** @hide */
+    public EventOutputParcel(@NonNull EventOutput value) {
+        this(value.getEventLogRecord());
+    }
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventOutputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    /**
+     * Creates a new EventOutputParcel.
+     *
+     * @param eventLogRecord
+     *   An {@link EventLogRecord} to be written to the EVENTS table, if not null. Each
+     *   {@link EventLogRecord} is associated with a row in an existing {@link RequestLogRecord} that
+     *   has been written to the REQUESTS table.
+     */
+    @DataClass.Generated.Member
+    public EventOutputParcel(
+            @Nullable EventLogRecord eventLogRecord) {
+        this.mEventLogRecord = eventLogRecord;
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * An {@link EventLogRecord} to be written to the EVENTS table, if not null. Each
+     * {@link EventLogRecord} is associated with a row in an existing {@link RequestLogRecord} that
+     * has been written to the REQUESTS table.
+     */
+    @DataClass.Generated.Member
+    public @Nullable EventLogRecord getEventLogRecord() {
+        return mEventLogRecord;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@android.annotation.NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        byte flg = 0;
+        if (mEventLogRecord != null) flg |= 0x1;
+        dest.writeByte(flg);
+        if (mEventLogRecord != null) dest.writeTypedObject(mEventLogRecord, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ EventOutputParcel(@android.annotation.NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        byte flg = in.readByte();
+        EventLogRecord eventLogRecord = (flg & 0x1) == 0 ? null : (EventLogRecord) in.readTypedObject(EventLogRecord.CREATOR);
+
+        this.mEventLogRecord = eventLogRecord;
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @android.annotation.NonNull Parcelable.Creator<EventOutputParcel> CREATOR
+            = new Parcelable.Creator<EventOutputParcel>() {
+        @Override
+        public EventOutputParcel[] newArray(int size) {
+            return new EventOutputParcel[size];
+        }
+
+        @Override
+        public EventOutputParcel createFromParcel(@android.annotation.NonNull android.os.Parcel in) {
+            return new EventOutputParcel(in);
+        }
+    };
+
+    @DataClass.Generated(
+            time = 1698864082503L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/EventOutputParcel.java",
+            inputSignatures = " @android.annotation.Nullable android.adservices.ondevicepersonalization.EventLogRecord mEventLogRecord\nclass EventOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genBuilder=false)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.aidl b/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.aidl
deleted file mode 100644
index 14f342f9..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.aidl
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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 android.adservices.ondevicepersonalization;
-
-parcelable ExecuteInput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
index b7fbdb8b..825b25a2 100644
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java
@@ -20,7 +20,6 @@ import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ON
 
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
-import android.os.Parcelable;
 import android.os.PersistableBundle;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
@@ -31,8 +30,8 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
  *
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
-@DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
-public final class ExecuteInput implements Parcelable {
+@DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
+public final class ExecuteInput {
     /**
      * The package name of the calling app.
      */
@@ -44,6 +43,10 @@ public final class ExecuteInput implements Parcelable {
      */
     @NonNull PersistableBundle mAppParams = PersistableBundle.EMPTY;
 
+    /** @hide */
+    public ExecuteInput(@NonNull ExecuteInputParcel parcel) {
+        this(parcel.getAppPackageName(), parcel.getAppParams());
+    }
 
 
 
@@ -60,8 +63,18 @@ public final class ExecuteInput implements Parcelable {
     //@formatter:off
 
 
+    /**
+     * Creates a new ExecuteInput.
+     *
+     * @param appPackageName
+     *   The package name of the calling app.
+     * @param appParams
+     *   The parameters provided by the app to the {@link IsolatedService}. The service
+     *   defines the expected keys in this {@link PersistableBundle}.
+     * @hide
+     */
     @DataClass.Generated.Member
-    /* package-private */ ExecuteInput(
+    public ExecuteInput(
             @NonNull String appPackageName,
             @NonNull PersistableBundle appParams) {
         this.mAppPackageName = appPackageName;
@@ -120,123 +133,11 @@ public final class ExecuteInput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        dest.writeString(mAppPackageName);
-        dest.writeTypedObject(mAppParams, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ ExecuteInput(@NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        String appPackageName = in.readString();
-        PersistableBundle appParams = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
-
-        this.mAppPackageName = appPackageName;
-        AnnotationValidations.validate(
-                NonNull.class, null, mAppPackageName);
-        this.mAppParams = appParams;
-        AnnotationValidations.validate(
-                NonNull.class, null, mAppParams);
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<ExecuteInput> CREATOR
-            = new Parcelable.Creator<ExecuteInput>() {
-        @Override
-        public ExecuteInput[] newArray(int size) {
-            return new ExecuteInput[size];
-        }
-
-        @Override
-        public ExecuteInput createFromParcel(@NonNull android.os.Parcel in) {
-            return new ExecuteInput(in);
-        }
-    };
-
-    /**
-     * A builder for {@link ExecuteInput}
-     * @hide
-     */
-    @SuppressWarnings("WeakerAccess")
-    @DataClass.Generated.Member
-    public static final class Builder {
-
-        private @NonNull String mAppPackageName;
-        private @NonNull PersistableBundle mAppParams;
-
-        private long mBuilderFieldsSet = 0L;
-
-        public Builder() {
-        }
-
-        /**
-         * The package name of the calling app.
-         */
-        @DataClass.Generated.Member
-        public @NonNull Builder setAppPackageName(@NonNull String value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x1;
-            mAppPackageName = value;
-            return this;
-        }
-
-        /**
-         * The parameters provided by the app to the {@link IsolatedService}. The service
-         * defines the expected keys in this {@link PersistableBundle}.
-         */
-        @DataClass.Generated.Member
-        public @NonNull Builder setAppParams(@NonNull PersistableBundle value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x2;
-            mAppParams = value;
-            return this;
-        }
-
-        /** Builds the instance. This builder should not be touched after calling this! */
-        public @NonNull ExecuteInput build() {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x4; // Mark builder used
-
-            if ((mBuilderFieldsSet & 0x1) == 0) {
-                mAppPackageName = "";
-            }
-            if ((mBuilderFieldsSet & 0x2) == 0) {
-                mAppParams = PersistableBundle.EMPTY;
-            }
-            ExecuteInput o = new ExecuteInput(
-                    mAppPackageName,
-                    mAppParams);
-            return o;
-        }
-
-        private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x4) != 0) {
-                throw new IllegalStateException(
-                        "This Builder should not be reused. Use a new Builder instance instead");
-            }
-        }
-    }
-
     @DataClass.Generated(
-            time = 1697051118497L,
+            time = 1698879004862L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteInput.java",
-            inputSignatures = " @android.annotation.NonNull java.lang.String mAppPackageName\n @android.annotation.NonNull android.os.PersistableBundle mAppParams\nclass ExecuteInput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genHiddenBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = " @android.annotation.NonNull java.lang.String mAppPackageName\n @android.annotation.NonNull android.os.PersistableBundle mAppParams\nclass ExecuteInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=false, genHiddenConstructor=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteInputParcel.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteInputParcel.java
new file mode 100644
index 00000000..44c60f30
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteInputParcel.java
@@ -0,0 +1,213 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+
+import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+/**
+ * Parcelable version of {@link ExecuteInput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genHiddenBuilder = true)
+public final class ExecuteInputParcel implements Parcelable {
+    /**
+     * The package name of the calling app.
+     */
+    @NonNull String mAppPackageName = "";
+
+    /**
+     * The parameters provided by the app to the {@link IsolatedService}. The service
+     * defines the expected keys in this {@link PersistableBundle}.
+     */
+    @NonNull PersistableBundle mAppParams = PersistableBundle.EMPTY;
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteInputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    @DataClass.Generated.Member
+    /* package-private */ ExecuteInputParcel(
+            @NonNull String appPackageName,
+            @NonNull PersistableBundle appParams) {
+        this.mAppPackageName = appPackageName;
+        AnnotationValidations.validate(
+                NonNull.class, null, mAppPackageName);
+        this.mAppParams = appParams;
+        AnnotationValidations.validate(
+                NonNull.class, null, mAppParams);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The package name of the calling app.
+     */
+    @DataClass.Generated.Member
+    public @NonNull String getAppPackageName() {
+        return mAppPackageName;
+    }
+
+    /**
+     * The parameters provided by the app to the {@link IsolatedService}. The service
+     * defines the expected keys in this {@link PersistableBundle}.
+     */
+    @DataClass.Generated.Member
+    public @NonNull PersistableBundle getAppParams() {
+        return mAppParams;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        dest.writeString(mAppPackageName);
+        dest.writeTypedObject(mAppParams, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ ExecuteInputParcel(@NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        String appPackageName = in.readString();
+        PersistableBundle appParams = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
+
+        this.mAppPackageName = appPackageName;
+        AnnotationValidations.validate(
+                NonNull.class, null, mAppPackageName);
+        this.mAppParams = appParams;
+        AnnotationValidations.validate(
+                NonNull.class, null, mAppParams);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @NonNull Parcelable.Creator<ExecuteInputParcel> CREATOR
+            = new Parcelable.Creator<ExecuteInputParcel>() {
+        @Override
+        public ExecuteInputParcel[] newArray(int size) {
+            return new ExecuteInputParcel[size];
+        }
+
+        @Override
+        public ExecuteInputParcel createFromParcel(@NonNull android.os.Parcel in) {
+            return new ExecuteInputParcel(in);
+        }
+    };
+
+    /**
+     * A builder for {@link ExecuteInputParcel}
+     * @hide
+     */
+    @SuppressWarnings("WeakerAccess")
+    @DataClass.Generated.Member
+    public static final class Builder {
+
+        private @NonNull String mAppPackageName;
+        private @NonNull PersistableBundle mAppParams;
+
+        private long mBuilderFieldsSet = 0L;
+
+        public Builder() {
+        }
+
+        /**
+         * The package name of the calling app.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setAppPackageName(@NonNull String value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x1;
+            mAppPackageName = value;
+            return this;
+        }
+
+        /**
+         * The parameters provided by the app to the {@link IsolatedService}. The service
+         * defines the expected keys in this {@link PersistableBundle}.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setAppParams(@NonNull PersistableBundle value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x2;
+            mAppParams = value;
+            return this;
+        }
+
+        /** Builds the instance. This builder should not be touched after calling this! */
+        public @NonNull ExecuteInputParcel build() {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x4; // Mark builder used
+
+            if ((mBuilderFieldsSet & 0x1) == 0) {
+                mAppPackageName = "";
+            }
+            if ((mBuilderFieldsSet & 0x2) == 0) {
+                mAppParams = PersistableBundle.EMPTY;
+            }
+            ExecuteInputParcel o = new ExecuteInputParcel(
+                    mAppPackageName,
+                    mAppParams);
+            return o;
+        }
+
+        private void checkNotUsed() {
+            if ((mBuilderFieldsSet & 0x4) != 0) {
+                throw new IllegalStateException(
+                        "This Builder should not be reused. Use a new Builder instance instead");
+            }
+        }
+    }
+
+    @DataClass.Generated(
+            time = 1698868678877L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteInputParcel.java",
+            inputSignatures = " @android.annotation.NonNull java.lang.String mAppPackageName\n @android.annotation.NonNull android.os.PersistableBundle mAppParams\nclass ExecuteInputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genHiddenBuilder=true)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.aidl b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.aidl
deleted file mode 100644
index 9746ecd6..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.aidl
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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 android.adservices.ondevicepersonalization;
-
-parcelable ExecuteOutput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
index ad7c5d2e..b6ee67dd 100644
--- a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java
@@ -21,7 +21,6 @@ import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ON
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Parcelable;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
@@ -38,7 +37,7 @@ import java.util.List;
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
-public final class ExecuteOutput implements Parcelable {
+public final class ExecuteOutput {
     /**
      * Persistent data to be written to the REQUESTS table after
      * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}
@@ -160,63 +159,6 @@ public final class ExecuteOutput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        byte flg = 0;
-        if (mRequestLogRecord != null) flg |= 0x1;
-        dest.writeByte(flg);
-        if (mRequestLogRecord != null) dest.writeTypedObject(mRequestLogRecord, flags);
-        dest.writeParcelableList(mRenderingConfigs, flags);
-        dest.writeParcelableList(mEventLogRecords, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ ExecuteOutput(@NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        byte flg = in.readByte();
-        RequestLogRecord requestLogRecord = (flg & 0x1) == 0 ? null : (RequestLogRecord) in.readTypedObject(RequestLogRecord.CREATOR);
-        List<RenderingConfig> renderingConfigs = new java.util.ArrayList<>();
-        in.readParcelableList(renderingConfigs, RenderingConfig.class.getClassLoader());
-        List<EventLogRecord> eventLogRecords = new java.util.ArrayList<>();
-        in.readParcelableList(eventLogRecords, EventLogRecord.class.getClassLoader());
-
-        this.mRequestLogRecord = requestLogRecord;
-        this.mRenderingConfigs = renderingConfigs;
-        AnnotationValidations.validate(
-                NonNull.class, null, mRenderingConfigs);
-        this.mEventLogRecords = eventLogRecords;
-        AnnotationValidations.validate(
-                NonNull.class, null, mEventLogRecords);
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<ExecuteOutput> CREATOR
-            = new Parcelable.Creator<ExecuteOutput>() {
-        @Override
-        public ExecuteOutput[] newArray(int size) {
-            return new ExecuteOutput[size];
-        }
-
-        @Override
-        public ExecuteOutput createFromParcel(@NonNull android.os.Parcel in) {
-            return new ExecuteOutput(in);
-        }
-    };
-
     /**
      * A builder for {@link ExecuteOutput}
      */
@@ -284,10 +226,7 @@ public final class ExecuteOutput implements Parcelable {
             return this;
         }
 
-        /**
-         * @see #setEventLogRecords
-         * @hide
-         */
+        /** @see #setEventLogRecords @hide */
         @DataClass.Generated.Member
         public @NonNull Builder addEventLogRecord(@NonNull EventLogRecord value) {
             if (mEventLogRecords == null) setEventLogRecords(new java.util.ArrayList<>());
@@ -325,10 +264,10 @@ public final class ExecuteOutput implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1697132452641L,
+            time = 1698880049845L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteOutput.java",
-            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"renderingConfig\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.RenderingConfig> mRenderingConfigs\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"eventLogRecord\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.EventLogRecord> mEventLogRecords\nclass ExecuteOutput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"renderingConfig\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.RenderingConfig> mRenderingConfigs\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"eventLogRecord\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.EventLogRecord> mEventLogRecords\nclass ExecuteOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/ExecuteOutputParcel.java b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutputParcel.java
new file mode 100644
index 00000000..df23d7cf
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/ExecuteOutputParcel.java
@@ -0,0 +1,216 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Parcelable;
+
+import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Parcelable version of {@link ExecuteOutput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genBuilder = false)
+public final class ExecuteOutputParcel implements Parcelable {
+    /**
+     * Persistent data to be written to the REQUESTS table after
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}
+     * completes. If null, no persistent data will be written.
+     */
+    @Nullable private RequestLogRecord mRequestLogRecord = null;
+
+    /**
+     * A list of {@link RenderingConfig} objects, one per slot specified in the request from the
+     * calling app. The calling app and the service must agree on the expected size of this list.
+     */
+    @DataClass.PluralOf("renderingConfig")
+    @NonNull private List<RenderingConfig> mRenderingConfigs = Collections.emptyList();
+
+    /**
+     * A list of {@link EventLogRecord}. Writes events to the EVENTS table and associates
+     * them with requests with the specified corresponding {@link RequestLogRecord} from
+     * {@link EventLogRecord#getRequestLogRecord()}.
+     * If the event does not contain a {@link RequestLogRecord} emitted by this package, the
+     * EventLogRecord is not written.
+     *
+     * @hide
+     */
+    @DataClass.PluralOf("eventLogRecord")
+    @NonNull private List<EventLogRecord> mEventLogRecords = Collections.emptyList();
+
+    /** @hide */
+    public ExecuteOutputParcel(@NonNull ExecuteOutput value) {
+        this(value.getRequestLogRecord(), value.getRenderingConfigs(), value.getEventLogRecords());
+    }
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteOutputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    /**
+     * Creates a new ExecuteOutputParcel.
+     *
+     * @param requestLogRecord
+     *   Persistent data to be written to the REQUESTS table after
+     *   {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}
+     *   completes. If null, no persistent data will be written.
+     * @param renderingConfigs
+     *   A list of {@link RenderingConfig} objects, one per slot specified in the request from the
+     *   calling app. The calling app and the service must agree on the expected size of this list.
+     * @param eventLogRecords
+     *   A list of {@link EventLogRecord}. Writes events to the EVENTS table and associates
+     *   them with requests with the specified corresponding {@link RequestLogRecord} from
+     *   {@link EventLogRecord#getRequestLogRecord()}.
+     *   If the event does not contain a {@link RequestLogRecord} emitted by this package, the
+     *   EventLogRecord is not written.
+     */
+    @DataClass.Generated.Member
+    public ExecuteOutputParcel(
+            @Nullable RequestLogRecord requestLogRecord,
+            @NonNull List<RenderingConfig> renderingConfigs,
+            @NonNull List<EventLogRecord> eventLogRecords) {
+        this.mRequestLogRecord = requestLogRecord;
+        this.mRenderingConfigs = renderingConfigs;
+        AnnotationValidations.validate(
+                NonNull.class, null, mRenderingConfigs);
+        this.mEventLogRecords = eventLogRecords;
+        AnnotationValidations.validate(
+                NonNull.class, null, mEventLogRecords);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * Persistent data to be written to the REQUESTS table after
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}
+     * completes. If null, no persistent data will be written.
+     */
+    @DataClass.Generated.Member
+    public @Nullable RequestLogRecord getRequestLogRecord() {
+        return mRequestLogRecord;
+    }
+
+    /**
+     * A list of {@link RenderingConfig} objects, one per slot specified in the request from the
+     * calling app. The calling app and the service must agree on the expected size of this list.
+     */
+    @DataClass.Generated.Member
+    public @NonNull List<RenderingConfig> getRenderingConfigs() {
+        return mRenderingConfigs;
+    }
+
+    /**
+     * A list of {@link EventLogRecord}. Writes events to the EVENTS table and associates
+     * them with requests with the specified corresponding {@link RequestLogRecord} from
+     * {@link EventLogRecord#getRequestLogRecord()}.
+     * If the event does not contain a {@link RequestLogRecord} emitted by this package, the
+     * EventLogRecord is not written.
+     *
+     * @hide
+     */
+    @DataClass.Generated.Member
+    public @NonNull List<EventLogRecord> getEventLogRecords() {
+        return mEventLogRecords;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        byte flg = 0;
+        if (mRequestLogRecord != null) flg |= 0x1;
+        dest.writeByte(flg);
+        if (mRequestLogRecord != null) dest.writeTypedObject(mRequestLogRecord, flags);
+        dest.writeParcelableList(mRenderingConfigs, flags);
+        dest.writeParcelableList(mEventLogRecords, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ ExecuteOutputParcel(@NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        byte flg = in.readByte();
+        RequestLogRecord requestLogRecord = (flg & 0x1) == 0 ? null : (RequestLogRecord) in.readTypedObject(RequestLogRecord.CREATOR);
+        List<RenderingConfig> renderingConfigs = new java.util.ArrayList<>();
+        in.readParcelableList(renderingConfigs, RenderingConfig.class.getClassLoader());
+        List<EventLogRecord> eventLogRecords = new java.util.ArrayList<>();
+        in.readParcelableList(eventLogRecords, EventLogRecord.class.getClassLoader());
+
+        this.mRequestLogRecord = requestLogRecord;
+        this.mRenderingConfigs = renderingConfigs;
+        AnnotationValidations.validate(
+                NonNull.class, null, mRenderingConfigs);
+        this.mEventLogRecords = eventLogRecords;
+        AnnotationValidations.validate(
+                NonNull.class, null, mEventLogRecords);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @NonNull Parcelable.Creator<ExecuteOutputParcel> CREATOR
+            = new Parcelable.Creator<ExecuteOutputParcel>() {
+        @Override
+        public ExecuteOutputParcel[] newArray(int size) {
+            return new ExecuteOutputParcel[size];
+        }
+
+        @Override
+        public ExecuteOutputParcel createFromParcel(@NonNull android.os.Parcel in) {
+            return new ExecuteOutputParcel(in);
+        }
+    };
+
+    @DataClass.Generated(
+            time = 1698864579986L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/ExecuteOutputParcel.java",
+            inputSignatures = "private @android.annotation.Nullable android.adservices.ondevicepersonalization.RequestLogRecord mRequestLogRecord\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"renderingConfig\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.RenderingConfig> mRenderingConfigs\nprivate @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"eventLogRecord\") @android.annotation.NonNull java.util.List<android.adservices.ondevicepersonalization.EventLogRecord> mEventLogRecords\nclass ExecuteOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genBuilder=false)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
index e5dee4c5..6057b2c1 100644
--- a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
@@ -16,54 +16,49 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
 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(FederatedComputeScheduler.Params, FederatedComputeInput)}
- *
- * @hide
- */
+/** The input data for {@link FederatedComputeScheduler#schedule}. */
 @DataClass(genBuilder = true, genEqualsHashCode = true)
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public final class FederatedComputeInput {
     // TODO(b/300461799): add federated compute server document.
     /**
      * Population refers to a collection of devices that specific task groups can run on. It should
-     * match task plan configured at remote federated computation server.
+     * match task plan configured at remote federated compute server.
      */
     @NonNull private String mPopulationName = "";
 
-
-
     // Code below generated by codegen v1.0.23.
     //
     // DO NOT MODIFY!
     // CHECKSTYLE:OFF Generated code
     //
     // To regenerate run:
-    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java
     //
     // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
     //   Settings > Editor > Code Style > Formatter Control
-    //@formatter:off
-
+    // @formatter:off
 
     @DataClass.Generated.Member
-    /* package-private */ FederatedComputeInput(
-            @NonNull String populationName) {
+    /* package-private */ FederatedComputeInput(@NonNull String populationName) {
         this.mPopulationName = populationName;
-        AnnotationValidations.validate(
-                NonNull.class, null, mPopulationName);
+        AnnotationValidations.validate(NonNull.class, null, mPopulationName);
 
         // onConstructed(); // You can define this method to get a callback
     }
 
     /**
      * Population refers to a collection of devices that specific task groups can run on. It should
-     * match task plan configured at remote federated computation server.
+     * match task plan configured at remote federated compute server.
      */
     @DataClass.Generated.Member
     public @NonNull String getPopulationName() {
@@ -82,8 +77,7 @@ public final class FederatedComputeInput {
         @SuppressWarnings("unchecked")
         FederatedComputeInput that = (FederatedComputeInput) o;
         //noinspection PointlessBooleanExpression
-        return true
-                && java.util.Objects.equals(mPopulationName, that.mPopulationName);
+        return true && java.util.Objects.equals(mPopulationName, that.mPopulationName);
     }
 
     @Override
@@ -97,9 +91,8 @@ public final class FederatedComputeInput {
         return _hash;
     }
 
-    /**
-     * A builder for {@link FederatedComputeInput}
-     */
+    /** A builder for {@link FederatedComputeInput} */
+    @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
     @SuppressWarnings("WeakerAccess")
     @DataClass.Generated.Member
     public static final class Builder {
@@ -108,13 +101,9 @@ public final class FederatedComputeInput {
 
         private long mBuilderFieldsSet = 0L;
 
-        public Builder() {
-        }
+        public Builder() {}
 
-        /**
-         * Population refers to a collection of devices that specific task groups can run on. It should
-         * match task plan configured at remote federated computation server.
-         */
+        /** Setter for {@link #getPopulationName}. */
         @DataClass.Generated.Member
         public @NonNull Builder setPopulationName(@NonNull String value) {
             checkNotUsed();
@@ -131,8 +120,7 @@ public final class FederatedComputeInput {
             if ((mBuilderFieldsSet & 0x1) == 0) {
                 mPopulationName = "";
             }
-            FederatedComputeInput o = new FederatedComputeInput(
-                    mPopulationName);
+            FederatedComputeInput o = new FederatedComputeInput(mPopulationName);
             return o;
         }
 
@@ -147,13 +135,14 @@ public final class FederatedComputeInput {
     @DataClass.Generated(
             time = 1697578140247L,
             codegenVersion = "1.0.23",
-            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java",
-            inputSignatures = "private @android.annotation.NonNull java.lang.String mPopulationName\nclass FederatedComputeInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            sourceFile =
+                    "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/FederatedComputeInput.java",
+            inputSignatures =
+                    "private @android.annotation.NonNull java.lang.String mPopulationName\nclass FederatedComputeInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
-
-    //@formatter:on
+    // @formatter:on
     // End of generated code
 
 }
diff --git a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
index 4077e7e8..6e6f2bef 100644
--- a/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
+++ b/framework/java/android/adservices/ondevicepersonalization/FederatedComputeScheduler.java
@@ -16,8 +16,11 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
 import android.adservices.ondevicepersonalization.aidl.IFederatedComputeCallback;
 import android.adservices.ondevicepersonalization.aidl.IFederatedComputeService;
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.WorkerThread;
 import android.federatedcompute.common.TrainingOptions;
@@ -28,10 +31,10 @@ import com.android.ondevicepersonalization.internal.util.LoggerFactory;
 import java.util.concurrent.CountDownLatch;
 
 /**
- * Handles scheduling Federated Learning and Federated Analytics jobs.
- *
- * @hide
+ * Handles scheduling federated compute jobs. See {@link
+ * IsolatedService#getFederatedComputeScheduler}.
  */
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class FederatedComputeScheduler {
     private static final String TAG = FederatedComputeScheduler.class.getSimpleName();
     private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
@@ -44,14 +47,15 @@ public class FederatedComputeScheduler {
     }
 
     // TODO(b/300461799): add federated compute server document.
+    // TODO(b/269665435): add sample code snippet.
     /**
-     * Schedule a federated computation job.
+     * Schedules a federated compute job. In {@link IsolatedService#onRequest}, the app can call
+     * {@link IsolatedService#getFederatedComputeScheduler} to pass scheduler when construct {@link
+     * IsolatedWorker}.
      *
      * @param params parameters related to job scheduling.
-     * @param input the configuration of the federated computation. It should be consistent with
-     *     the federated computation server setup.
-     * @throws IllegalArgumentException caused by caller supplied invalid input argument.
-     * @throws IllegalStateException caused by an internal failure of FederatedComputeScheduler.
+     * @param input the configuration of the federated compute. It should be consistent with the
+     *     federated compute server setup.
      */
     @WorkerThread
     public void schedule(@NonNull Params params, @NonNull FederatedComputeInput input) {
@@ -94,7 +98,10 @@ public class FederatedComputeScheduler {
     }
 
     /**
-     * Cancel a federated computation job with input training params.
+     * Cancels a federated compute job with input training params. In {@link
+     * IsolatedService#onRequest}, the app can call {@link
+     * IsolatedService#getFederatedComputeScheduler} to pass scheduler when construct {@link
+     * IsolatedWorker}.
      *
      * @param populationName population name of the job that caller wants to cancel
      * @throws IllegalStateException caused by an internal failure of FederatedComputeScheduler.
@@ -141,6 +148,7 @@ public class FederatedComputeScheduler {
     }
 
     /** The parameters related to job scheduling. */
+    @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
     public static class Params {
         /**
          * If training interval is scheduled for recurrent tasks, the earliest time this task could
diff --git a/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java b/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
index f9fcc59b..5a8fd11a 100644
--- a/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
+++ b/framework/java/android/adservices/ondevicepersonalization/IsolatedService.java
@@ -40,6 +40,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 // TODO(b/289102463): Add a link to the public ODP developer documentation.
 /**
@@ -144,8 +145,6 @@ public abstract class IsolatedService extends Service {
      *     The methods in the returned {@link LogReader} are blocking operations and
      *     should be called from a worker thread and not the main thread or a binder thread.
      * @see #onRequest(RequestToken)
-     *
-     * @hide
      */
     @NonNull
     public final LogReader getLogReader(@NonNull RequestToken requestToken) {
@@ -181,14 +180,13 @@ public abstract class IsolatedService extends Service {
 
     /**
      * Returns an {@link FederatedComputeScheduler} for the current request. The {@link
-     * FederatedComputeScheduler} can be used to schedule and cancel federated computation jobs. The
-     * federated computation includes federated learning and federated analytic jobs.
+     * FederatedComputeScheduler} can be used to schedule and cancel federated computation jobs.
+     * The federated computation includes federated learning and federated analytic jobs.
      *
      * @param requestToken an opaque token that identifies the current request to the service.
      * @return An {@link FederatedComputeScheduler} that returns a federated computation job
      *     scheduler.
      * @see #onRequest(RequestToken)
-     * @hide
      */
     @NonNull
     public final FederatedComputeScheduler getFederatedComputeScheduler(
@@ -210,9 +208,9 @@ public abstract class IsolatedService extends Service {
 
             if (operationCode == Constants.OP_EXECUTE) {
 
-                ExecuteInput input =
-                        Objects.requireNonNull(
-                                params.getParcelable(Constants.EXTRA_INPUT, ExecuteInput.class));
+                ExecuteInputParcel inputParcel = Objects.requireNonNull(
+                        params.getParcelable(Constants.EXTRA_INPUT, ExecuteInputParcel.class));
+                ExecuteInput input = new ExecuteInput(inputParcel);
                 Objects.requireNonNull(input.getAppPackageName());
                 IDataAccessService binder =
                         IDataAccessService.Stub.asInterface(
@@ -229,8 +227,10 @@ public abstract class IsolatedService extends Service {
                 UserData userData = params.getParcelable(Constants.EXTRA_USER_DATA, UserData.class);
                 RequestToken requestToken = new RequestToken(binder, fcBinder, userData);
                 IsolatedWorker implCallback = IsolatedService.this.onRequest(requestToken);
-                implCallback.onExecute(input, new WrappedCallback<ExecuteOutput>(
-                        resultCallback, requestToken));
+                implCallback.onExecute(
+                        input,
+                        new WrappedCallback<ExecuteOutput, ExecuteOutputParcel>(
+                            resultCallback, requestToken, v -> new ExecuteOutputParcel(v)));
 
             } else if (operationCode == Constants.OP_DOWNLOAD) {
 
@@ -274,14 +274,19 @@ public abstract class IsolatedService extends Service {
                 RequestToken requestToken = new RequestToken(binder, fcBinder, userData);
                 IsolatedWorker implCallback = IsolatedService.this.onRequest(requestToken);
                 implCallback.onDownloadCompleted(
-                        input, new WrappedCallback<DownloadCompletedOutput>(
-                                resultCallback, requestToken));
+                        input,
+                        new WrappedCallback<DownloadCompletedOutput,
+                                DownloadCompletedOutputParcel>(
+                                resultCallback,
+                                requestToken,
+                                v -> new DownloadCompletedOutputParcel(v)));
 
             } else if (operationCode == Constants.OP_RENDER) {
 
-                RenderInput input =
-                        Objects.requireNonNull(
-                                params.getParcelable(Constants.EXTRA_INPUT, RenderInput.class));
+                RenderInputParcel inputParcel =
+                        Objects.requireNonNull(params.getParcelable(
+                                Constants.EXTRA_INPUT, RenderInputParcel.class));
+                RenderInput input = new RenderInput(inputParcel);
                 Objects.requireNonNull(input.getRenderingConfig());
                 IDataAccessService binder =
                         IDataAccessService.Stub.asInterface(
@@ -291,15 +296,16 @@ public abstract class IsolatedService extends Service {
                 Objects.requireNonNull(binder);
                 RequestToken requestToken = new RequestToken(binder, null, null);
                 IsolatedWorker implCallback = IsolatedService.this.onRequest(requestToken);
-                implCallback.onRender(input, new WrappedCallback<RenderOutput>(
-                        resultCallback, requestToken));
+                implCallback.onRender(input, new WrappedCallback<RenderOutput, RenderOutputParcel>(
+                        resultCallback, requestToken, v -> new RenderOutputParcel(v)));
 
             } else if (operationCode == Constants.OP_WEB_VIEW_EVENT) {
 
-                EventInput input =
+                EventInputParcel inputParcel =
                         Objects.requireNonNull(
                                 params.getParcelable(
-                                        Constants.EXTRA_INPUT, EventInput.class));
+                                        Constants.EXTRA_INPUT, EventInputParcel.class));
+                EventInput input = new EventInput(inputParcel);
                 IDataAccessService binder =
                         IDataAccessService.Stub.asInterface(
                                 Objects.requireNonNull(
@@ -309,13 +315,15 @@ public abstract class IsolatedService extends Service {
                 RequestToken requestToken = new RequestToken(binder, null, userData);
                 IsolatedWorker implCallback = IsolatedService.this.onRequest(requestToken);
                 implCallback.onEvent(
-                        input, new WrappedCallback<EventOutput>(resultCallback, requestToken));
+                        input, new WrappedCallback<EventOutput, EventOutputParcel>(
+                            resultCallback, requestToken, v -> new EventOutputParcel(v)));
 
             } else if (operationCode == Constants.OP_TRAINING_EXAMPLE) {
-                TrainingExampleInput input =
+                TrainingExamplesInputParcel inputParcel =
                         Objects.requireNonNull(
                                 params.getParcelable(
-                                        Constants.EXTRA_INPUT, TrainingExampleInput.class));
+                                        Constants.EXTRA_INPUT, TrainingExamplesInputParcel.class));
+                TrainingExamplesInput input = new TrainingExamplesInput(inputParcel);
                 IDataAccessService binder =
                         IDataAccessService.Stub.asInterface(
                                 Objects.requireNonNull(
@@ -325,12 +333,14 @@ public abstract class IsolatedService extends Service {
                 UserData userData = params.getParcelable(Constants.EXTRA_USER_DATA, UserData.class);
                 RequestToken requestToken = new RequestToken(binder, null, userData);
                 IsolatedWorker implCallback = IsolatedService.this.onRequest(requestToken);
-                implCallback.onTrainingExample(
-                        input, new Consumer<TrainingExampleOutput>() {
+                implCallback.onTrainingExamples(
+                        input,
+                        new Consumer<TrainingExamplesOutput>() {
                             @Override
-                            public void accept(TrainingExampleOutput result) {
-                                long elapsedTimeMillis = SystemClock.elapsedRealtime()
-                                        - requestToken.getStartTimeMillis();
+                            public void accept(TrainingExamplesOutput result) {
+                                long elapsedTimeMillis =
+                                        SystemClock.elapsedRealtime()
+                                                - requestToken.getStartTimeMillis();
                                 if (result == null) {
                                     try {
                                         resultCallback.onError(Constants.STATUS_INTERNAL_ERROR);
@@ -338,8 +348,8 @@ public abstract class IsolatedService extends Service {
                                         sLogger.w(TAG + ": Callback failed.", e);
                                     }
                                 } else {
-                                    TrainingExampleOutputParcel parcelResult =
-                                            new TrainingExampleOutputParcel.Builder()
+                                    TrainingExamplesOutputParcel parcelResult =
+                                            new TrainingExamplesOutputParcel.Builder()
                                                     .setTrainingExamples(
                                                             new ByteArrayParceledListSlice(
                                                                     result.getTrainingExamples()))
@@ -349,10 +359,11 @@ public abstract class IsolatedService extends Service {
                                                     .build();
                                     Bundle bundle = new Bundle();
                                     bundle.putParcelable(Constants.EXTRA_RESULT, parcelResult);
-                                    bundle.putParcelable(Constants.EXTRA_CALLEE_METADATA,
+                                    bundle.putParcelable(
+                                            Constants.EXTRA_CALLEE_METADATA,
                                             new CalleeMetadata.Builder()
-                                                .setElapsedTimeMillis(elapsedTimeMillis)
-                                                .build());
+                                                    .setElapsedTimeMillis(elapsedTimeMillis)
+                                                    .build());
                                     try {
                                         resultCallback.onSuccess(bundle);
                                     } catch (RemoteException e) {
@@ -367,13 +378,18 @@ public abstract class IsolatedService extends Service {
         }
     }
 
-    private static class WrappedCallback<T extends Parcelable> implements Consumer<T> {
+    private static class WrappedCallback<T, U extends Parcelable> implements Consumer<T> {
         @NonNull private final IIsolatedServiceCallback mCallback;
         @NonNull private final RequestToken mRequestToken;
+        @NonNull private final Function<T, U> mConverter;
 
-        WrappedCallback(IIsolatedServiceCallback callback, RequestToken requestToken) {
+        WrappedCallback(
+                IIsolatedServiceCallback callback,
+                RequestToken requestToken,
+                Function<T, U> converter) {
             mCallback = Objects.requireNonNull(callback);
             mRequestToken = Objects.requireNonNull(requestToken);
+            mConverter = Objects.requireNonNull(converter);
         }
 
         @Override
@@ -388,7 +404,8 @@ public abstract class IsolatedService extends Service {
                 }
             } else {
                 Bundle bundle = new Bundle();
-                bundle.putParcelable(Constants.EXTRA_RESULT, result);
+                U wrappedResult = mConverter.apply(result);
+                bundle.putParcelable(Constants.EXTRA_RESULT, wrappedResult);
                 bundle.putParcelable(Constants.EXTRA_CALLEE_METADATA,
                         new CalleeMetadata.Builder()
                             .setElapsedTimeMillis(elapsedTimeMillis)
diff --git a/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java b/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
index 39adfc36..11a651eb 100644
--- a/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
+++ b/framework/java/android/adservices/ondevicepersonalization/IsolatedWorker.java
@@ -106,7 +106,6 @@ public interface IsolatedWorker {
      *     an error. If called with <code>null</code>, no data is written to the EVENTS table.
      *     <p>If this method throws a {@link RuntimeException}, no data is written to the EVENTS
      *     table.
-     * @hide
      */
     default void onEvent(
             @NonNull EventInput input, @NonNull Consumer<EventOutput> consumer) {
@@ -114,18 +113,18 @@ public interface IsolatedWorker {
     }
 
     /**
-     * Generate a single training example used for federated computation job.
+     * Generate a list of training examples used for federated compute job. The platform will call
+     * this function when a federated compute job starts. The federated compute job is scheduled by
+     * an app through {@link FederatedComputeScheduler#schedule}.
      *
      * @param input The parameters needed to generate the training example.
      * @param consumer Callback that receives the result. Should be called with <code>null</code> on
      *     an error. If called with <code>null</code>, no training examples is produced for this
-     *     training session. <p>If this method throws a {@link RuntimeException}, no training
-     *     examples are produced for this training session.
-     * @hide
+     *     training session.
      */
-    default void onTrainingExample(
-            @NonNull TrainingExampleInput input,
-            @NonNull Consumer<TrainingExampleOutput> consumer) {
+    default void onTrainingExamples(
+            @NonNull TrainingExamplesInput input,
+            @NonNull Consumer<TrainingExamplesOutput> consumer) {
         consumer.accept(null);
     }
 }
diff --git a/framework/java/android/adservices/ondevicepersonalization/LogReader.java b/framework/java/android/adservices/ondevicepersonalization/LogReader.java
index 7d54dd55..bf25efcf 100644
--- a/framework/java/android/adservices/ondevicepersonalization/LogReader.java
+++ b/framework/java/android/adservices/ondevicepersonalization/LogReader.java
@@ -16,8 +16,11 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
 import android.adservices.ondevicepersonalization.aidl.IDataAccessService;
 import android.adservices.ondevicepersonalization.aidl.IDataAccessServiceCallback;
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.WorkerThread;
 import android.os.Bundle;
@@ -27,6 +30,7 @@ import android.os.RemoteException;
 import com.android.ondevicepersonalization.internal.util.LoggerFactory;
 import com.android.ondevicepersonalization.internal.util.OdpParceledListSlice;
 
+import java.time.Instant;
 import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.ArrayBlockingQueue;
@@ -39,8 +43,8 @@ import java.util.concurrent.BlockingQueue;
  *
  * @see IsolatedService#getLogReader(RequestToken)
  *
- * @hide
  */
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 public class LogReader {
     private static final String TAG = "LogReader";
     private static final LoggerFactory.Logger sLogger = LoggerFactory.getLogger();
@@ -60,7 +64,10 @@ public class LogReader {
      */
     @WorkerThread
     @NonNull
-    public List<RequestLogRecord> getRequests(long startTimeMillis, long endTimeMillis) {
+    public List<RequestLogRecord> getRequests(
+            @NonNull Instant startTime, @NonNull Instant endTime) {
+        long startTimeMillis = startTime.toEpochMilli();
+        long endTimeMillis = endTime.toEpochMilli();
         if (endTimeMillis <= startTimeMillis) {
             throw new IllegalArgumentException(
                     "endTimeMillis must be greater than startTimeMillis");
@@ -82,7 +89,10 @@ public class LogReader {
      */
     @WorkerThread
     @NonNull
-    public List<EventLogRecord> getJoinedEvents(long startTimeMillis, long endTimeMillis) {
+    public List<EventLogRecord> getJoinedEvents(
+            @NonNull Instant startTime, @NonNull Instant endTime) {
+        long startTimeMillis = startTime.toEpochMilli();
+        long endTimeMillis = endTime.toEpochMilli();
         if (endTimeMillis <= startTimeMillis) {
             throw new IllegalArgumentException(
                     "endTimeMillis must be greater than startTimeMillis");
diff --git a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
index 5c0e9b5c..ab772f3a 100644
--- a/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
+++ b/framework/java/android/adservices/ondevicepersonalization/OnDevicePersonalizationManager.java
@@ -111,7 +111,6 @@ public class OnDevicePersonalizationManager {
      *     package is not installed or does not have a valid ODP manifest.
      *     Returns {@link ClassNotFoundException} if the handler class is not found.
      *     Returns an {@link OnDevicePersonalizationException} if execution of the handler fails.
-     * @hide
      */
     public void execute(
             @NonNull ComponentName handler,
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderInput.aidl b/framework/java/android/adservices/ondevicepersonalization/RenderInput.aidl
deleted file mode 100644
index dc5d618a..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/RenderInput.aidl
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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 android.adservices.ondevicepersonalization;
-
-parcelable RenderInput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderInput.java b/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
index a9a41df7..0c5824de 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderInput.java
@@ -19,8 +19,8 @@ package android.adservices.ondevicepersonalization;
 import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
 
 import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Parcelable;
 
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
@@ -30,8 +30,8 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
  *
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
-@DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
-public final class RenderInput implements Parcelable {
+@DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
+public final class RenderInput {
     /** The width of the slot. */
     private int mWidth = 0;
 
@@ -50,6 +50,12 @@ public final class RenderInput implements Parcelable {
      */
     @Nullable RenderingConfig mRenderingConfig = null;
 
+    /** @hide */
+    public RenderInput(@NonNull RenderInputParcel parcel) {
+        this(parcel.getWidth(), parcel.getHeight(), parcel.getRenderingConfigIndex(),
+                parcel.getRenderingConfig());
+    }
+
 
 
     // Code below generated by codegen v1.0.23.
@@ -65,8 +71,23 @@ public final class RenderInput implements Parcelable {
     //@formatter:off
 
 
+    /**
+     * Creates a new RenderInput.
+     *
+     * @param width
+     *   The width of the slot.
+     * @param height
+     *   The height of the slot.
+     * @param renderingConfigIndex
+     *   The index of the {@link RenderingConfig} in {@link ExecuteOutput} that this render
+     *   request is for.
+     * @param renderingConfig
+     *   A {@link RenderingConfig} within an {@link ExecuteOutput} that was returned by
+     *   {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     * @hide
+     */
     @DataClass.Generated.Member
-    /* package-private */ RenderInput(
+    public RenderInput(
             int width,
             int height,
             int renderingConfigIndex,
@@ -146,162 +167,11 @@ public final class RenderInput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@android.annotation.NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        byte flg = 0;
-        if (mRenderingConfig != null) flg |= 0x8;
-        dest.writeByte(flg);
-        dest.writeInt(mWidth);
-        dest.writeInt(mHeight);
-        dest.writeInt(mRenderingConfigIndex);
-        if (mRenderingConfig != null) dest.writeTypedObject(mRenderingConfig, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ RenderInput(@android.annotation.NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        byte flg = in.readByte();
-        int width = in.readInt();
-        int height = in.readInt();
-        int renderingConfigIndex = in.readInt();
-        RenderingConfig renderingConfig = (flg & 0x8) == 0 ? null : (RenderingConfig) in.readTypedObject(RenderingConfig.CREATOR);
-
-        this.mWidth = width;
-        this.mHeight = height;
-        this.mRenderingConfigIndex = renderingConfigIndex;
-        this.mRenderingConfig = renderingConfig;
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @android.annotation.NonNull Parcelable.Creator<RenderInput> CREATOR
-            = new Parcelable.Creator<RenderInput>() {
-        @Override
-        public RenderInput[] newArray(int size) {
-            return new RenderInput[size];
-        }
-
-        @Override
-        public RenderInput createFromParcel(@android.annotation.NonNull android.os.Parcel in) {
-            return new RenderInput(in);
-        }
-    };
-
-    /**
-     * A builder for {@link RenderInput}
-     * @hide
-     */
-    @SuppressWarnings("WeakerAccess")
-    @DataClass.Generated.Member
-    public static final class Builder {
-
-        private int mWidth;
-        private int mHeight;
-        private int mRenderingConfigIndex;
-        private @Nullable RenderingConfig mRenderingConfig;
-
-        private long mBuilderFieldsSet = 0L;
-
-        public Builder() {
-        }
-
-        /**
-         * The width of the slot.
-         */
-        @DataClass.Generated.Member
-        public @android.annotation.NonNull Builder setWidth(int value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x1;
-            mWidth = value;
-            return this;
-        }
-
-        /**
-         * The height of the slot.
-         */
-        @DataClass.Generated.Member
-        public @android.annotation.NonNull Builder setHeight(int value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x2;
-            mHeight = value;
-            return this;
-        }
-
-        /**
-         * The index of the {@link RenderingConfig} in {@link ExecuteOutput} that this render
-         * request is for.
-         */
-        @DataClass.Generated.Member
-        public @android.annotation.NonNull Builder setRenderingConfigIndex(int value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x4;
-            mRenderingConfigIndex = value;
-            return this;
-        }
-
-        /**
-         * A {@link RenderingConfig} within an {@link ExecuteOutput} that was returned by
-         * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
-         */
-        @DataClass.Generated.Member
-        public @android.annotation.NonNull Builder setRenderingConfig(@android.annotation.NonNull RenderingConfig value) {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x8;
-            mRenderingConfig = value;
-            return this;
-        }
-
-        /** Builds the instance. This builder should not be touched after calling this! */
-        public @android.annotation.NonNull RenderInput build() {
-            checkNotUsed();
-            mBuilderFieldsSet |= 0x10; // Mark builder used
-
-            if ((mBuilderFieldsSet & 0x1) == 0) {
-                mWidth = 0;
-            }
-            if ((mBuilderFieldsSet & 0x2) == 0) {
-                mHeight = 0;
-            }
-            if ((mBuilderFieldsSet & 0x4) == 0) {
-                mRenderingConfigIndex = 0;
-            }
-            if ((mBuilderFieldsSet & 0x8) == 0) {
-                mRenderingConfig = null;
-            }
-            RenderInput o = new RenderInput(
-                    mWidth,
-                    mHeight,
-                    mRenderingConfigIndex,
-                    mRenderingConfig);
-            return o;
-        }
-
-        private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x10) != 0) {
-                throw new IllegalStateException(
-                        "This Builder should not be reused. Use a new Builder instance instead");
-            }
-        }
-    }
-
     @DataClass.Generated(
-            time = 1697063555914L,
+            time = 1698879032836L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderInput.java",
-            inputSignatures = "private  int mWidth\nprivate  int mHeight\nprivate  int mRenderingConfigIndex\n @android.annotation.Nullable android.adservices.ondevicepersonalization.RenderingConfig mRenderingConfig\nclass RenderInput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genHiddenBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = "private  int mWidth\nprivate  int mHeight\nprivate  int mRenderingConfigIndex\n @android.annotation.Nullable android.adservices.ondevicepersonalization.RenderingConfig mRenderingConfig\nclass RenderInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=false, genHiddenConstructor=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderInputParcel.java b/framework/java/android/adservices/ondevicepersonalization/RenderInputParcel.java
new file mode 100644
index 00000000..bf2af756
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderInputParcel.java
@@ -0,0 +1,274 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import android.annotation.Nullable;
+import android.os.Parcelable;
+
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+/**
+ * Parcelable version of {@link RenderInput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genHiddenBuilder = true)
+public final class RenderInputParcel implements Parcelable {
+    /** The width of the slot. */
+    private int mWidth = 0;
+
+    /** The height of the slot. */
+    private int mHeight = 0;
+
+    /**
+     * The index of the {@link RenderingConfig} in {@link ExecuteOutput} that this render
+     * request is for.
+     */
+    private int mRenderingConfigIndex = 0;
+
+    /**
+     * A {@link RenderingConfig} within an {@link ExecuteOutput} that was returned by
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     */
+    @Nullable RenderingConfig mRenderingConfig = null;
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderInputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    @DataClass.Generated.Member
+    /* package-private */ RenderInputParcel(
+            int width,
+            int height,
+            int renderingConfigIndex,
+            @Nullable RenderingConfig renderingConfig) {
+        this.mWidth = width;
+        this.mHeight = height;
+        this.mRenderingConfigIndex = renderingConfigIndex;
+        this.mRenderingConfig = renderingConfig;
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The width of the slot.
+     */
+    @DataClass.Generated.Member
+    public int getWidth() {
+        return mWidth;
+    }
+
+    /**
+     * The height of the slot.
+     */
+    @DataClass.Generated.Member
+    public int getHeight() {
+        return mHeight;
+    }
+
+    /**
+     * The index of the {@link RenderingConfig} in {@link ExecuteOutput} that this render
+     * request is for.
+     */
+    @DataClass.Generated.Member
+    public int getRenderingConfigIndex() {
+        return mRenderingConfigIndex;
+    }
+
+    /**
+     * A {@link RenderingConfig} within an {@link ExecuteOutput} that was returned by
+     * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+     */
+    @DataClass.Generated.Member
+    public @Nullable RenderingConfig getRenderingConfig() {
+        return mRenderingConfig;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@android.annotation.NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        byte flg = 0;
+        if (mRenderingConfig != null) flg |= 0x8;
+        dest.writeByte(flg);
+        dest.writeInt(mWidth);
+        dest.writeInt(mHeight);
+        dest.writeInt(mRenderingConfigIndex);
+        if (mRenderingConfig != null) dest.writeTypedObject(mRenderingConfig, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ RenderInputParcel(@android.annotation.NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        byte flg = in.readByte();
+        int width = in.readInt();
+        int height = in.readInt();
+        int renderingConfigIndex = in.readInt();
+        RenderingConfig renderingConfig = (flg & 0x8) == 0 ? null : (RenderingConfig) in.readTypedObject(RenderingConfig.CREATOR);
+
+        this.mWidth = width;
+        this.mHeight = height;
+        this.mRenderingConfigIndex = renderingConfigIndex;
+        this.mRenderingConfig = renderingConfig;
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @android.annotation.NonNull Parcelable.Creator<RenderInputParcel> CREATOR
+            = new Parcelable.Creator<RenderInputParcel>() {
+        @Override
+        public RenderInputParcel[] newArray(int size) {
+            return new RenderInputParcel[size];
+        }
+
+        @Override
+        public RenderInputParcel createFromParcel(@android.annotation.NonNull android.os.Parcel in) {
+            return new RenderInputParcel(in);
+        }
+    };
+
+    /**
+     * A builder for {@link RenderInputParcel}
+     * @hide
+     */
+    @SuppressWarnings("WeakerAccess")
+    @DataClass.Generated.Member
+    public static final class Builder {
+
+        private int mWidth;
+        private int mHeight;
+        private int mRenderingConfigIndex;
+        private @Nullable RenderingConfig mRenderingConfig;
+
+        private long mBuilderFieldsSet = 0L;
+
+        public Builder() {
+        }
+
+        /**
+         * The width of the slot.
+         */
+        @DataClass.Generated.Member
+        public @android.annotation.NonNull Builder setWidth(int value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x1;
+            mWidth = value;
+            return this;
+        }
+
+        /**
+         * The height of the slot.
+         */
+        @DataClass.Generated.Member
+        public @android.annotation.NonNull Builder setHeight(int value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x2;
+            mHeight = value;
+            return this;
+        }
+
+        /**
+         * The index of the {@link RenderingConfig} in {@link ExecuteOutput} that this render
+         * request is for.
+         */
+        @DataClass.Generated.Member
+        public @android.annotation.NonNull Builder setRenderingConfigIndex(int value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x4;
+            mRenderingConfigIndex = value;
+            return this;
+        }
+
+        /**
+         * A {@link RenderingConfig} within an {@link ExecuteOutput} that was returned by
+         * {@link IsolatedWorker#onExecute(ExecuteInput, java.util.function.Consumer)}.
+         */
+        @DataClass.Generated.Member
+        public @android.annotation.NonNull Builder setRenderingConfig(@android.annotation.NonNull RenderingConfig value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x8;
+            mRenderingConfig = value;
+            return this;
+        }
+
+        /** Builds the instance. This builder should not be touched after calling this! */
+        public @android.annotation.NonNull RenderInputParcel build() {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x10; // Mark builder used
+
+            if ((mBuilderFieldsSet & 0x1) == 0) {
+                mWidth = 0;
+            }
+            if ((mBuilderFieldsSet & 0x2) == 0) {
+                mHeight = 0;
+            }
+            if ((mBuilderFieldsSet & 0x4) == 0) {
+                mRenderingConfigIndex = 0;
+            }
+            if ((mBuilderFieldsSet & 0x8) == 0) {
+                mRenderingConfig = null;
+            }
+            RenderInputParcel o = new RenderInputParcel(
+                    mWidth,
+                    mHeight,
+                    mRenderingConfigIndex,
+                    mRenderingConfig);
+            return o;
+        }
+
+        private void checkNotUsed() {
+            if ((mBuilderFieldsSet & 0x10) != 0) {
+                throw new IllegalStateException(
+                        "This Builder should not be reused. Use a new Builder instance instead");
+            }
+        }
+    }
+
+    @DataClass.Generated(
+            time = 1698872925083L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderInputParcel.java",
+            inputSignatures = "private  int mWidth\nprivate  int mHeight\nprivate  int mRenderingConfigIndex\n @android.annotation.Nullable android.adservices.ondevicepersonalization.RenderingConfig mRenderingConfig\nclass RenderInputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genHiddenBuilder=true)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.aidl b/framework/java/android/adservices/ondevicepersonalization/RenderOutput.aidl
deleted file mode 100644
index 42c26fff..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.aidl
+++ /dev/null
@@ -1,19 +0,0 @@
-/*
- * 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 android.adservices.ondevicepersonalization;
-
-parcelable RenderOutput;
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java b/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
index 3b7c8dc1..33c8c8b3 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java
@@ -21,7 +21,6 @@ import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ON
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Parcelable;
 import android.os.PersistableBundle;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
@@ -34,7 +33,7 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
  */
 @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
-public final class RenderOutput implements Parcelable {
+public final class RenderOutput {
     /**
      * The HTML content to be rendered in a webview. If this is null, the ODP service
      * generates HTML from the data in {@link #getTemplateId()} and {@link #getTemplateParams()}
@@ -146,60 +145,6 @@ public final class RenderOutput implements Parcelable {
         return _hash;
     }
 
-    @Override
-    @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
-        // You can override field parcelling by defining methods like:
-        // void parcelFieldName(Parcel dest, int flags) { ... }
-
-        byte flg = 0;
-        if (mContent != null) flg |= 0x1;
-        if (mTemplateId != null) flg |= 0x2;
-        dest.writeByte(flg);
-        if (mContent != null) dest.writeString(mContent);
-        if (mTemplateId != null) dest.writeString(mTemplateId);
-        dest.writeTypedObject(mTemplateParams, flags);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int describeContents() { return 0; }
-
-    /** @hide */
-    @SuppressWarnings({"unchecked", "RedundantCast"})
-    @DataClass.Generated.Member
-    /* package-private */ RenderOutput(@NonNull android.os.Parcel in) {
-        // You can override field unparcelling by defining methods like:
-        // static FieldType unparcelFieldName(Parcel in) { ... }
-
-        byte flg = in.readByte();
-        String content = (flg & 0x1) == 0 ? null : in.readString();
-        String templateId = (flg & 0x2) == 0 ? null : in.readString();
-        PersistableBundle templateParams = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
-
-        this.mContent = content;
-        this.mTemplateId = templateId;
-        this.mTemplateParams = templateParams;
-        AnnotationValidations.validate(
-                NonNull.class, null, mTemplateParams);
-
-        // onConstructed(); // You can define this method to get a callback
-    }
-
-    @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<RenderOutput> CREATOR
-            = new Parcelable.Creator<RenderOutput>() {
-        @Override
-        public RenderOutput[] newArray(int size) {
-            return new RenderOutput[size];
-        }
-
-        @Override
-        public RenderOutput createFromParcel(@NonNull android.os.Parcel in) {
-            return new RenderOutput(in);
-        }
-    };
-
     /**
      * A builder for {@link RenderOutput}
      */
@@ -285,10 +230,10 @@ public final class RenderOutput implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1697132582732L,
+            time = 1698880093023L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderOutput.java",
-            inputSignatures = "private @android.annotation.Nullable java.lang.String mContent\nprivate @android.annotation.Nullable java.lang.String mTemplateId\nprivate @android.annotation.NonNull android.os.PersistableBundle mTemplateParams\nclass RenderOutput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = "private @android.annotation.Nullable java.lang.String mContent\nprivate @android.annotation.Nullable java.lang.String mTemplateId\nprivate @android.annotation.NonNull android.os.PersistableBundle mTemplateParams\nclass RenderOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/RenderOutputParcel.java b/framework/java/android/adservices/ondevicepersonalization/RenderOutputParcel.java
new file mode 100644
index 00000000..afd4cfaf
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/RenderOutputParcel.java
@@ -0,0 +1,197 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Parcelable;
+import android.os.PersistableBundle;
+
+import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
+import com.android.ondevicepersonalization.internal.util.DataClass;
+
+/**
+ * Parcelable version of {@link RenderOutput}.
+ * @hide
+ */
+@DataClass(genAidl = false, genBuilder = false)
+public final class RenderOutputParcel implements Parcelable {
+    /**
+     * The HTML content to be rendered in a webview. If this is null, the ODP service
+     * generates HTML from the data in {@link #getTemplateId()} and {@link #getTemplateParams()}
+     * as described below.
+     */
+    @Nullable private String mContent = null;
+
+    /**
+     * A key in the REMOTE_DATA {@link IsolatedService#getRemoteData(RequestToken)} table that
+     * points to an <a href="velocity.apache.org">Apache Velocity</a> template. This is ignored if
+     * {@link #getContent()} is not null.
+     */
+    @Nullable private String mTemplateId = null;
+
+    /**
+     * The parameters to be populated in the template from {@link #getTemplateId()}. This is
+     * ignored if {@link #getContent()} is not null.
+     */
+    @NonNull private PersistableBundle mTemplateParams = PersistableBundle.EMPTY;
+
+    /** @hide */
+    public RenderOutputParcel(@NonNull RenderOutput value) {
+        this(value.getContent(), value.getTemplateId(), value.getTemplateParams());
+    }
+
+
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderOutputParcel.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    //@formatter:off
+
+
+    /**
+     * Creates a new RenderOutputParcel.
+     *
+     * @param content
+     *   The HTML content to be rendered in a webview. If this is null, the ODP service
+     *   generates HTML from the data in {@link #getTemplateId()} and {@link #getTemplateParams()}
+     *   as described below.
+     * @param templateId
+     *   A key in the REMOTE_DATA {@link IsolatedService#getRemoteData(RequestToken)} table that
+     *   points to an <a href="velocity.apache.org">Apache Velocity</a> template. This is ignored if
+     *   {@link #getContent()} is not null.
+     * @param templateParams
+     *   The parameters to be populated in the template from {@link #getTemplateId()}. This is
+     *   ignored if {@link #getContent()} is not null.
+     */
+    @DataClass.Generated.Member
+    public RenderOutputParcel(
+            @Nullable String content,
+            @Nullable String templateId,
+            @NonNull PersistableBundle templateParams) {
+        this.mContent = content;
+        this.mTemplateId = templateId;
+        this.mTemplateParams = templateParams;
+        AnnotationValidations.validate(
+                NonNull.class, null, mTemplateParams);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The HTML content to be rendered in a webview. If this is null, the ODP service
+     * generates HTML from the data in {@link #getTemplateId()} and {@link #getTemplateParams()}
+     * as described below.
+     */
+    @DataClass.Generated.Member
+    public @Nullable String getContent() {
+        return mContent;
+    }
+
+    /**
+     * A key in the REMOTE_DATA {@link IsolatedService#getRemoteData(RequestToken)} table that
+     * points to an <a href="velocity.apache.org">Apache Velocity</a> template. This is ignored if
+     * {@link #getContent()} is not null.
+     */
+    @DataClass.Generated.Member
+    public @Nullable String getTemplateId() {
+        return mTemplateId;
+    }
+
+    /**
+     * The parameters to be populated in the template from {@link #getTemplateId()}. This is
+     * ignored if {@link #getContent()} is not null.
+     */
+    @DataClass.Generated.Member
+    public @NonNull PersistableBundle getTemplateParams() {
+        return mTemplateParams;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+        // You can override field parcelling by defining methods like:
+        // void parcelFieldName(Parcel dest, int flags) { ... }
+
+        byte flg = 0;
+        if (mContent != null) flg |= 0x1;
+        if (mTemplateId != null) flg |= 0x2;
+        dest.writeByte(flg);
+        if (mContent != null) dest.writeString(mContent);
+        if (mTemplateId != null) dest.writeString(mTemplateId);
+        dest.writeTypedObject(mTemplateParams, flags);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int describeContents() { return 0; }
+
+    /** @hide */
+    @SuppressWarnings({"unchecked", "RedundantCast"})
+    @DataClass.Generated.Member
+    /* package-private */ RenderOutputParcel(@NonNull android.os.Parcel in) {
+        // You can override field unparcelling by defining methods like:
+        // static FieldType unparcelFieldName(Parcel in) { ... }
+
+        byte flg = in.readByte();
+        String content = (flg & 0x1) == 0 ? null : in.readString();
+        String templateId = (flg & 0x2) == 0 ? null : in.readString();
+        PersistableBundle templateParams = (PersistableBundle) in.readTypedObject(PersistableBundle.CREATOR);
+
+        this.mContent = content;
+        this.mTemplateId = templateId;
+        this.mTemplateParams = templateParams;
+        AnnotationValidations.validate(
+                NonNull.class, null, mTemplateParams);
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    @DataClass.Generated.Member
+    public static final @NonNull Parcelable.Creator<RenderOutputParcel> CREATOR
+            = new Parcelable.Creator<RenderOutputParcel>() {
+        @Override
+        public RenderOutputParcel[] newArray(int size) {
+            return new RenderOutputParcel[size];
+        }
+
+        @Override
+        public RenderOutputParcel createFromParcel(@NonNull android.os.Parcel in) {
+            return new RenderOutputParcel(in);
+        }
+    };
+
+    @DataClass.Generated(
+            time = 1698864341247L,
+            codegenVersion = "1.0.23",
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RenderOutputParcel.java",
+            inputSignatures = "private @android.annotation.Nullable java.lang.String mContent\nprivate @android.annotation.Nullable java.lang.String mTemplateId\nprivate @android.annotation.NonNull android.os.PersistableBundle mTemplateParams\nclass RenderOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genBuilder=false)")
+    @Deprecated
+    private void __metadata() {}
+
+
+    //@formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java b/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
index 8cd16764..618462e8 100644
--- a/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
+++ b/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java
@@ -26,6 +26,7 @@ import android.os.Parcelable;
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
+import java.time.Instant;
 import java.util.Collections;
 import java.util.List;
 
@@ -60,6 +61,13 @@ public final class RequestLogRecord implements Parcelable {
      */
     private long mTimeMillis = 0;
 
+    /**
+     * Returns the timestamp of this record.
+     */
+    @NonNull public Instant getTime() {
+        return Instant.ofEpochMilli(getTimeMillis());
+    }
+
     abstract static class BaseBuilder {
         /**
          * @hide
@@ -298,10 +306,10 @@ public final class RequestLogRecord implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1696978492795L,
+            time = 1698962042612L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/RequestLogRecord.java",
-            inputSignatures = " @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"row\") @android.annotation.NonNull java.util.List<android.content.ContentValues> mRows\nprivate  long mRequestId\nprivate  long mTimeMillis\nclass RequestLogRecord extends java.lang.Object implements [android.os.Parcelable]\npublic abstract  android.adservices.ondevicepersonalization.RequestLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)\npublic abstract  android.adservices.ondevicepersonalization.RequestLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []")
+            inputSignatures = " @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"row\") @android.annotation.NonNull java.util.List<android.content.ContentValues> mRows\nprivate  long mRequestId\nprivate  long mTimeMillis\npublic @android.annotation.NonNull java.time.Instant getTime()\nclass RequestLogRecord extends java.lang.Object implements [android.os.Parcelable]\npublic abstract  android.adservices.ondevicepersonalization.RequestLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)\npublic abstract  android.adservices.ondevicepersonalization.RequestLogRecord.Builder setTimeMillis(long)\nclass BaseBuilder extends java.lang.Object implements []")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.aidl b/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.aidl
deleted file mode 100644
index 699e9c59..00000000
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.aidl
+++ /dev/null
@@ -1,19 +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.adservices.ondevicepersonalization;
-
-parcelable TrainingExampleOutputParcel;
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java
new file mode 100644
index 00000000..bd8ca99e
--- /dev/null
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInput.java
@@ -0,0 +1,168 @@
+/*
+ * 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 android.adservices.ondevicepersonalization;
+
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+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}. */
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
+@DataClass(genBuilder = false, genHiddenConstructor = true, genEqualsHashCode = true)
+public final class TrainingExamplesInput {
+    /**
+     * The name of the federated compute population. It should match the population name in {@link
+     * FederatedComputeInput#getPopulationName}.
+     */
+    @NonNull private String mPopulationName = "";
+
+    /**
+     * The name of the task within the population. It should match task plan configured at remote
+     * federated compute server. One population may have multiple tasks. The task name can be used
+     * to uniquely identify the job.
+     */
+    @NonNull private String mTaskName = "";
+
+    /**
+     * Token used to support the resumption of training. If client app wants to use resumption token
+     * to track what examples are already used in previous federated compute jobs, it need set
+     * {@link TrainingExamplesOutput.Builder#setResumptionTokens}, OnDevicePersonalization will
+     * store it and pass it here for generating new training examples.
+     */
+    @Nullable private byte[] mResumptionToken = null;
+
+    /** @hide */
+    public TrainingExamplesInput(@NonNull TrainingExamplesInputParcel parcel) {
+        this(parcel.getPopulationName(), parcel.getTaskName(), parcel.getResumptionToken());
+    }
+
+    // Code below generated by codegen v1.0.23.
+    //
+    // DO NOT MODIFY!
+    // CHECKSTYLE:OFF Generated code
+    //
+    // To regenerate run:
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java
+    //
+    // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
+    //   Settings > Editor > Code Style > Formatter Control
+    // @formatter:off
+
+    /**
+     * Creates a new TrainingExampleInput.
+     *
+     * @param populationName The name of the federated compute population.
+     * @param taskName The name of the task within the population. One population may have multiple
+     *     tasks. The task name can be used to uniquely identify the job.
+     * @param resumptionToken Token used to support the resumption of training.
+     * @hide
+     */
+    @DataClass.Generated.Member
+    public TrainingExamplesInput(
+            @NonNull String populationName,
+            @NonNull String taskName,
+            @Nullable byte[] resumptionToken) {
+        this.mPopulationName = populationName;
+        AnnotationValidations.validate(NonNull.class, null, mPopulationName);
+        this.mTaskName = taskName;
+        AnnotationValidations.validate(NonNull.class, null, mTaskName);
+        this.mResumptionToken = resumptionToken;
+
+        // onConstructed(); // You can define this method to get a callback
+    }
+
+    /**
+     * The name of the federated compute population. It should match the population name in {@link
+     * FederatedComputeInput#getPopulationName}.
+     */
+    @DataClass.Generated.Member
+    public @NonNull String getPopulationName() {
+        return mPopulationName;
+    }
+
+    /**
+     * The name of the task within the population. It should match task plan configured at remote
+     * federated compute server. One population may have multiple tasks. The task name can be used
+     * to uniquely identify the job.
+     */
+    @DataClass.Generated.Member
+    public @NonNull String getTaskName() {
+        return mTaskName;
+    }
+
+    /**
+     * Token used to support the resumption of training. If client app wants to use resumption token
+     * to track what examples are already used in previous federated compute jobs, it need set
+     * {@link TrainingExamplesOutput.Builder#setResumptionTokens}, OnDevicePersonalization will
+     * store it and pass it here for generating new training examples.
+     */
+    @DataClass.Generated.Member
+    public @Nullable byte[] getResumptionToken() {
+        return mResumptionToken;
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public boolean equals(@Nullable Object o) {
+        // You can override field equality logic by defining either of the methods like:
+        // boolean fieldNameEquals(TrainingExampleInput other) { ... }
+        // boolean fieldNameEquals(FieldType otherValue) { ... }
+
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        @SuppressWarnings("unchecked")
+        TrainingExamplesInput that = (TrainingExamplesInput) o;
+        //noinspection PointlessBooleanExpression
+        return true
+                && java.util.Objects.equals(mPopulationName, that.mPopulationName)
+                && java.util.Objects.equals(mTaskName, that.mTaskName)
+                && java.util.Arrays.equals(mResumptionToken, that.mResumptionToken);
+    }
+
+    @Override
+    @DataClass.Generated.Member
+    public int hashCode() {
+        // You can override field hashCode logic by defining methods like:
+        // int fieldNameHashCode() { ... }
+
+        int _hash = 1;
+        _hash = 31 * _hash + java.util.Objects.hashCode(mPopulationName);
+        _hash = 31 * _hash + java.util.Objects.hashCode(mTaskName);
+        _hash = 31 * _hash + java.util.Arrays.hashCode(mResumptionToken);
+        return _hash;
+    }
+
+    @DataClass.Generated(
+            time = 1699394018457L,
+            codegenVersion = "1.0.23",
+            sourceFile =
+                    "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java",
+            inputSignatures =
+                    "private @android.annotation.NonNull java.lang.String mPopulationName\nprivate @android.annotation.NonNull java.lang.String mTaskName\nprivate @android.annotation.Nullable byte[] mResumptionToken\nclass TrainingExampleInput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=false, genHiddenConstructor=true, genEqualsHashCode=true)")
+    @Deprecated
+    private void __metadata() {}
+
+    // @formatter:on
+    // End of generated code
+
+}
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInputParcel.java
index 412f99aa..9e14b545 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesInputParcel.java
@@ -23,53 +23,47 @@ import android.os.Parcelable;
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
-import java.util.function.Consumer;
-
 /**
- * The input data for {@link IsolatedWorker#onTrainingExample(TrainingExampleInput, Consumer)}
+ * Parcelable version of {@link TrainingExamplesInput}
  *
  * @hide
  */
-@DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
-public final class TrainingExampleInput implements Parcelable {
+@DataClass(genAidl = false, genHiddenBuilder = true)
+public final class TrainingExamplesInputParcel implements Parcelable {
     /** The name of the federated compute population. */
     @NonNull private String mPopulationName = "";
 
     /**
-     * The name of the task within the population. One population may have multiple tasks.
-     * The task name can be used to uniquely identify the job.
+     * The name of the task within the population. One population may have multiple tasks. The task
+     * name can be used to uniquely identify the job.
      */
     @NonNull private String mTaskName = "";
 
     /** Token used to support the resumption of training. */
     @Nullable private byte[] mResumptionToken = null;
 
-
-
     // Code below generated by codegen v1.0.23.
     //
     // DO NOT MODIFY!
     // CHECKSTYLE:OFF Generated code
     //
     // To regenerate run:
-    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInputParcel.java
     //
     // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
     //   Settings > Editor > Code Style > Formatter Control
-    //@formatter:off
-
+    // @formatter:off
 
     @DataClass.Generated.Member
-    /* package-private */ TrainingExampleInput(
+    /* package-private */ TrainingExamplesInputParcel(
             @NonNull String populationName,
             @NonNull String taskName,
             @Nullable byte[] resumptionToken) {
         this.mPopulationName = populationName;
-        AnnotationValidations.validate(
-                NonNull.class, null, mPopulationName);
+        AnnotationValidations.validate(NonNull.class, null, mPopulationName);
         this.mTaskName = taskName;
-        AnnotationValidations.validate(
-                NonNull.class, null, mTaskName);
+        AnnotationValidations.validate(NonNull.class, null, mTaskName);
         this.mResumptionToken = resumptionToken;
 
         // onConstructed(); // You can define this method to get a callback
@@ -102,37 +96,6 @@ public final class TrainingExampleInput implements Parcelable {
 
     @Override
     @DataClass.Generated.Member
-    public boolean equals(@Nullable Object o) {
-        // You can override field equality logic by defining either of the methods like:
-        // boolean fieldNameEquals(TrainingExampleInput other) { ... }
-        // boolean fieldNameEquals(FieldType otherValue) { ... }
-
-        if (this == o) return true;
-        if (o == null || getClass() != o.getClass()) return false;
-        @SuppressWarnings("unchecked")
-        TrainingExampleInput that = (TrainingExampleInput) o;
-        //noinspection PointlessBooleanExpression
-        return true
-                && java.util.Objects.equals(mPopulationName, that.mPopulationName)
-                && java.util.Objects.equals(mTaskName, that.mTaskName)
-                && java.util.Arrays.equals(mResumptionToken, that.mResumptionToken);
-    }
-
-    @Override
-    @DataClass.Generated.Member
-    public int hashCode() {
-        // You can override field hashCode logic by defining methods like:
-        // int fieldNameHashCode() { ... }
-
-        int _hash = 1;
-        _hash = 31 * _hash + java.util.Objects.hashCode(mPopulationName);
-        _hash = 31 * _hash + java.util.Objects.hashCode(mTaskName);
-        _hash = 31 * _hash + java.util.Arrays.hashCode(mResumptionToken);
-        return _hash;
-    }
-
-    @Override
-    @DataClass.Generated.Member
     public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
         // You can override field parcelling by defining methods like:
         // void parcelFieldName(Parcel dest, int flags) { ... }
@@ -149,7 +112,7 @@ public final class TrainingExampleInput implements Parcelable {
     /** @hide */
     @SuppressWarnings({"unchecked", "RedundantCast"})
     @DataClass.Generated.Member
-    /* package-private */ TrainingExampleInput(@NonNull android.os.Parcel in) {
+    /* package-private */ TrainingExamplesInputParcel(@NonNull android.os.Parcel in) {
         // You can override field unparcelling by defining methods like:
         // static FieldType unparcelFieldName(Parcel in) { ... }
 
@@ -158,8 +121,7 @@ public final class TrainingExampleInput implements Parcelable {
         byte[] resumptionToken = in.createByteArray();
 
         this.mPopulationName = populationName;
-        AnnotationValidations.validate(
-                NonNull.class, null, mPopulationName);
+        AnnotationValidations.validate(NonNull.class, null, mPopulationName);
         this.mTaskName = taskName;
         AnnotationValidations.validate(
                 NonNull.class, null, mTaskName);
@@ -169,21 +131,22 @@ public final class TrainingExampleInput implements Parcelable {
     }
 
     @DataClass.Generated.Member
-    public static final @NonNull Parcelable.Creator<TrainingExampleInput> CREATOR
-            = new Parcelable.Creator<TrainingExampleInput>() {
-        @Override
-        public TrainingExampleInput[] newArray(int size) {
-            return new TrainingExampleInput[size];
-        }
-
-        @Override
-        public TrainingExampleInput createFromParcel(@NonNull android.os.Parcel in) {
-            return new TrainingExampleInput(in);
-        }
-    };
+    public static final @NonNull Parcelable.Creator<TrainingExamplesInputParcel> CREATOR =
+            new Parcelable.Creator<TrainingExamplesInputParcel>() {
+                @Override
+                public TrainingExamplesInputParcel[] newArray(int size) {
+                    return new TrainingExamplesInputParcel[size];
+                }
+
+                @Override
+                public TrainingExamplesInputParcel createFromParcel(@NonNull android.os.Parcel in) {
+                    return new TrainingExamplesInputParcel(in);
+                }
+            };
 
     /**
-     * A builder for {@link TrainingExampleInput}
+     * A builder for {@link TrainingExamplesInputParcel}
+     *
      * @hide
      */
     @SuppressWarnings("WeakerAccess")
@@ -234,7 +197,7 @@ public final class TrainingExampleInput implements Parcelable {
         }
 
         /** Builds the instance. This builder should not be touched after calling this! */
-        public @NonNull TrainingExampleInput build() {
+        public @NonNull TrainingExamplesInputParcel build() {
             checkNotUsed();
             mBuilderFieldsSet |= 0x8; // Mark builder used
 
@@ -247,10 +210,8 @@ public final class TrainingExampleInput implements Parcelable {
             if ((mBuilderFieldsSet & 0x4) == 0) {
                 mResumptionToken = null;
             }
-            TrainingExampleInput o = new TrainingExampleInput(
-                    mPopulationName,
-                    mTaskName,
-                    mResumptionToken);
+            TrainingExamplesInputParcel o =
+                    new TrainingExamplesInputParcel(mPopulationName, mTaskName, mResumptionToken);
             return o;
         }
 
@@ -263,10 +224,10 @@ public final class TrainingExampleInput implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1697577073626L,
+            time = 1699393441579L,
             codegenVersion = "1.0.23",
-            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInput.java",
-            inputSignatures = "private @android.annotation.NonNull java.lang.String mPopulationName\nprivate @android.annotation.NonNull java.lang.String mTaskName\nprivate @android.annotation.Nullable byte[] mResumptionToken\nclass TrainingExampleInput extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genHiddenBuilder=true, genEqualsHashCode=true)")
+            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleInputParcel.java",
+            inputSignatures = "private @android.annotation.NonNull java.lang.String mPopulationName\nprivate @android.annotation.NonNull java.lang.String mTaskName\nprivate @android.annotation.Nullable byte[] mResumptionToken\nclass TrainingExampleInputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genHiddenBuilder=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java
index 1a698686..deea25cb 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutput.java
@@ -16,6 +16,9 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 
 import com.android.internal.util.Preconditions;
@@ -24,18 +27,14 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 
 import java.util.Collections;
 import java.util.List;
-import java.util.function.Consumer;
 
-/**
- * The output data of {@link IsolatedWorker#onTrainingExample(TrainingExampleInput, Consumer)}
- *
- * @hide
- */
+/** The output data of {@link IsolatedWorker#onTrainingExamples} */
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genEqualsHashCode = true)
-public final class TrainingExampleOutput {
+public final class TrainingExamplesOutput {
     /**
-     * A list of training example byte arrays. The format is a binary serialized
-     * <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
+     * A list of training example byte arrays. The format is a binary serialized <a
+     * href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
      * tensorflow.Example</a> proto. The maximum allowed example size is 50KB.
      */
     @NonNull
@@ -43,52 +42,46 @@ public final class TrainingExampleOutput {
     private List<byte[]> mTrainingExamples = Collections.emptyList();
 
     /**
-     * A list of resumption token byte arrays corresponding to training examples. The last
-     * processed example's corresponding resumption token will be passed to
-     * {@link IsolatedWorker#onTrainingExample(TrainingExampleInput, Consumer)} to support
-     * resumption. The length of this list must match the length {@link #getTrainingExamples()}.
+     * A list of resumption token byte arrays corresponding to training examples. The last processed
+     * example's corresponding resumption token will be passed to {@link
+     * IsolatedWorker#onTrainingExamples} to support resumption. The length of this list must match
+     * the length {@link #getTrainingExamples}.
      */
     @NonNull
     @DataClass.PluralOf("resumptionToken")
     private List<byte[]> mResumptionTokens = Collections.emptyList();
 
-
     private void onConstructed() {
         Preconditions.checkArgument(mTrainingExamples.size() == mResumptionTokens.size());
     }
 
-
-
     // Code below generated by codegen v1.0.23.
     //
     // DO NOT MODIFY!
     // CHECKSTYLE:OFF Generated code
     //
     // To regenerate run:
-    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java
     //
     // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
     //   Settings > Editor > Code Style > Formatter Control
-    //@formatter:off
-
+    // @formatter:off
 
     @DataClass.Generated.Member
-    /* package-private */ TrainingExampleOutput(
-            @NonNull List<byte[]> trainingExamples,
-            @NonNull List<byte[]> resumptionTokens) {
+    /* package-private */ TrainingExamplesOutput(
+            @NonNull List<byte[]> trainingExamples, @NonNull List<byte[]> resumptionTokens) {
         this.mTrainingExamples = trainingExamples;
-        AnnotationValidations.validate(
-                NonNull.class, null, mTrainingExamples);
+        AnnotationValidations.validate(NonNull.class, null, mTrainingExamples);
         this.mResumptionTokens = resumptionTokens;
-        AnnotationValidations.validate(
-                NonNull.class, null, mResumptionTokens);
+        AnnotationValidations.validate(NonNull.class, null, mResumptionTokens);
 
         onConstructed();
     }
 
     /**
-     * A list of training example byte arrays. The format is a binary serialized
-     * <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
+     * A list of training example byte arrays. The format is a binary serialized <a
+     * href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
      * tensorflow.Example</a> proto. The maximum allowed example size is 50KB.
      */
     @DataClass.Generated.Member
@@ -97,10 +90,10 @@ public final class TrainingExampleOutput {
     }
 
     /**
-     * A list of resumption token byte arrays corresponding to training examples. The last
-     * processed example's corresponding resumption token will be passed to
-     * {@link IsolatedWorker#onTrainingExample(TrainingExampleInput, Consumer)} to support
-     * resumption. The length of this list must match the length {@link #getTrainingExamples()}.
+     * A list of resumption token byte arrays corresponding to training examples. The last processed
+     * example's corresponding resumption token will be passed to {@link
+     * IsolatedWorker#onTrainingExamples} to support resumption. The length of this list must match
+     * the length {@link #getTrainingExamples}.
      */
     @DataClass.Generated.Member
     public @NonNull List<byte[]> getResumptionTokens() {
@@ -117,7 +110,7 @@ public final class TrainingExampleOutput {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         @SuppressWarnings("unchecked")
-        TrainingExampleOutput that = (TrainingExampleOutput) o;
+        TrainingExamplesOutput that = (TrainingExamplesOutput) o;
         //noinspection PointlessBooleanExpression
         return true
                 && java.util.Objects.equals(mTrainingExamples, that.mTrainingExamples)
@@ -136,9 +129,8 @@ public final class TrainingExampleOutput {
         return _hash;
     }
 
-    /**
-     * A builder for {@link TrainingExampleOutput}
-     */
+    /** A builder for {@link TrainingExamplesOutput} */
+    @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
     @SuppressWarnings("WeakerAccess")
     @DataClass.Generated.Member
     public static final class Builder {
@@ -148,12 +140,11 @@ public final class TrainingExampleOutput {
 
         private long mBuilderFieldsSet = 0L;
 
-        public Builder() {
-        }
+        public Builder() {}
 
         /**
-         * A list of training example byte arrays. The format is a binary serialized
-         * <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
+         * A list of training example byte arrays. The format is a binary serialized <a
+         * href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/example/example.proto">
          * tensorflow.Example</a> proto. The maximum allowed example size is 50KB.
          */
         @DataClass.Generated.Member
@@ -164,7 +155,9 @@ public final class TrainingExampleOutput {
             return this;
         }
 
-        /** @see #setTrainingExamples */
+        /**
+         * @see #setTrainingExamples
+         */
         @DataClass.Generated.Member
         public @NonNull Builder addTrainingExample(@NonNull byte[] value) {
             if (mTrainingExamples == null) setTrainingExamples(new java.util.ArrayList<>());
@@ -174,9 +167,9 @@ public final class TrainingExampleOutput {
 
         /**
          * A list of resumption token byte arrays corresponding to training examples. The last
-         * processed example's corresponding resumption token will be passed to
-         * {@link IsolatedWorker#onTrainingExample(TrainingExampleInput, Consumer)} to support
-         * resumption. The length of this list must match the length {@link #getTrainingExamples()}.
+         * processed example's corresponding resumption token will be passed to {@link
+         * IsolatedWorker#onTrainingExamples} to support resumption. The length of this list must
+         * match the length {@link #getTrainingExamples}.
          */
         @DataClass.Generated.Member
         public @NonNull Builder setResumptionTokens(@NonNull List<byte[]> value) {
@@ -186,7 +179,9 @@ public final class TrainingExampleOutput {
             return this;
         }
 
-        /** @see #setResumptionTokens */
+        /**
+         * @see #setResumptionTokens
+         */
         @DataClass.Generated.Member
         public @NonNull Builder addResumptionToken(@NonNull byte[] value) {
             if (mResumptionTokens == null) setResumptionTokens(new java.util.ArrayList<>());
@@ -195,7 +190,7 @@ public final class TrainingExampleOutput {
         }
 
         /** Builds the instance. This builder should not be touched after calling this! */
-        public @NonNull TrainingExampleOutput build() {
+        public @NonNull TrainingExamplesOutput build() {
             checkNotUsed();
             mBuilderFieldsSet |= 0x4; // Mark builder used
 
@@ -205,9 +200,8 @@ public final class TrainingExampleOutput {
             if ((mBuilderFieldsSet & 0x2) == 0) {
                 mResumptionTokens = Collections.emptyList();
             }
-            TrainingExampleOutput o = new TrainingExampleOutput(
-                    mTrainingExamples,
-                    mResumptionTokens);
+            TrainingExamplesOutput o =
+                    new TrainingExamplesOutput(mTrainingExamples, mResumptionTokens);
             return o;
         }
 
@@ -222,13 +216,14 @@ public final class TrainingExampleOutput {
     @DataClass.Generated(
             time = 1697575854959L,
             codegenVersion = "1.0.23",
-            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java",
-            inputSignatures = "private @android.annotation.NonNull @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"trainingExample\") java.util.List<byte[]> mTrainingExamples\nprivate @android.annotation.NonNull @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"resumptionToken\") java.util.List<byte[]> mResumptionTokens\nprivate  void onConstructed()\nclass TrainingExampleOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
+            sourceFile =
+                    "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutput.java",
+            inputSignatures =
+                    "private @android.annotation.NonNull @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"trainingExample\") java.util.List<byte[]> mTrainingExamples\nprivate @android.annotation.NonNull @com.android.ondevicepersonalization.internal.util.DataClass.PluralOf(\"resumptionToken\") java.util.List<byte[]> mResumptionTokens\nprivate  void onConstructed()\nclass TrainingExampleOutput extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
-
-    //@formatter:on
+    // @formatter:on
     // End of generated code
 
 }
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.java b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutputParcel.java
index 8bbd518b..b620c098 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingExamplesOutputParcel.java
@@ -23,21 +23,17 @@ import com.android.ondevicepersonalization.internal.util.ByteArrayParceledListSl
 import com.android.ondevicepersonalization.internal.util.DataClass;
 
 /**
- * Parcelable version of {@link TrainingExampleOutput}
+ * Parcelable version of {@link TrainingExamplesOutput}
  *
  * @hide
  */
-@DataClass(genHiddenBuilder = true, genEqualsHashCode = true)
-public class TrainingExampleOutputParcel implements Parcelable {
+@DataClass(genAidl = false, genHiddenBuilder = true, genEqualsHashCode = true)
+public class TrainingExamplesOutputParcel implements Parcelable {
     /** List of training examples */
-    @Nullable
-    ByteArrayParceledListSlice mTrainingExamples = null;
+    @Nullable ByteArrayParceledListSlice mTrainingExamples = null;
 
     /** List of resumption tokens */
-    @Nullable
-    ByteArrayParceledListSlice mResumptionTokens = null;
-
-
+    @Nullable ByteArrayParceledListSlice mResumptionTokens = null;
 
     // Code below generated by codegen v1.0.23.
     //
@@ -45,15 +41,15 @@ public class TrainingExampleOutputParcel implements Parcelable {
     // CHECKSTYLE:OFF Generated code
     //
     // To regenerate run:
-    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.java
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.java
     //
     // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
     //   Settings > Editor > Code Style > Formatter Control
-    //@formatter:off
-
+    // @formatter:off
 
     @DataClass.Generated.Member
-    /* package-private */ TrainingExampleOutputParcel(
+    /* package-private */ TrainingExamplesOutputParcel(
             @Nullable ByteArrayParceledListSlice trainingExamples,
             @Nullable ByteArrayParceledListSlice resumptionTokens) {
         this.mTrainingExamples = trainingExamples;
@@ -88,7 +84,7 @@ public class TrainingExampleOutputParcel implements Parcelable {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         @SuppressWarnings("unchecked")
-        TrainingExampleOutputParcel that = (TrainingExampleOutputParcel) o;
+        TrainingExamplesOutputParcel that = (TrainingExamplesOutputParcel) o;
         //noinspection PointlessBooleanExpression
         return true
                 && java.util.Objects.equals(mTrainingExamples, that.mTrainingExamples)
@@ -128,13 +124,21 @@ public class TrainingExampleOutputParcel implements Parcelable {
     /** @hide */
     @SuppressWarnings({"unchecked", "RedundantCast"})
     @DataClass.Generated.Member
-    protected TrainingExampleOutputParcel(@android.annotation.NonNull android.os.Parcel in) {
+    protected TrainingExamplesOutputParcel(@android.annotation.NonNull android.os.Parcel in) {
         // You can override field unparcelling by defining methods like:
         // static FieldType unparcelFieldName(Parcel in) { ... }
 
         byte flg = in.readByte();
-        ByteArrayParceledListSlice trainingExamples = (flg & 0x1) == 0 ? null : (ByteArrayParceledListSlice) in.readTypedObject(ByteArrayParceledListSlice.CREATOR);
-        ByteArrayParceledListSlice resumptionTokens = (flg & 0x2) == 0 ? null : (ByteArrayParceledListSlice) in.readTypedObject(ByteArrayParceledListSlice.CREATOR);
+        ByteArrayParceledListSlice trainingExamples =
+                (flg & 0x1) == 0
+                        ? null
+                        : (ByteArrayParceledListSlice)
+                                in.readTypedObject(ByteArrayParceledListSlice.CREATOR);
+        ByteArrayParceledListSlice resumptionTokens =
+                (flg & 0x2) == 0
+                        ? null
+                        : (ByteArrayParceledListSlice)
+                                in.readTypedObject(ByteArrayParceledListSlice.CREATOR);
 
         this.mTrainingExamples = trainingExamples;
         this.mResumptionTokens = resumptionTokens;
@@ -143,21 +147,24 @@ public class TrainingExampleOutputParcel implements Parcelable {
     }
 
     @DataClass.Generated.Member
-    public static final @android.annotation.NonNull Parcelable.Creator<TrainingExampleOutputParcel> CREATOR
-            = new Parcelable.Creator<TrainingExampleOutputParcel>() {
-        @Override
-        public TrainingExampleOutputParcel[] newArray(int size) {
-            return new TrainingExampleOutputParcel[size];
-        }
-
-        @Override
-        public TrainingExampleOutputParcel createFromParcel(@android.annotation.NonNull android.os.Parcel in) {
-            return new TrainingExampleOutputParcel(in);
-        }
-    };
+    public static final @android.annotation.NonNull Parcelable.Creator<TrainingExamplesOutputParcel>
+            CREATOR =
+                    new Parcelable.Creator<TrainingExamplesOutputParcel>() {
+                        @Override
+                        public TrainingExamplesOutputParcel[] newArray(int size) {
+                            return new TrainingExamplesOutputParcel[size];
+                        }
+
+                        @Override
+                        public TrainingExamplesOutputParcel createFromParcel(
+                                @android.annotation.NonNull android.os.Parcel in) {
+                            return new TrainingExamplesOutputParcel(in);
+                        }
+                    };
 
     /**
-     * A builder for {@link TrainingExampleOutputParcel}
+     * A builder for {@link TrainingExamplesOutputParcel}
+     *
      * @hide
      */
     @SuppressWarnings("WeakerAccess")
@@ -195,7 +202,7 @@ public class TrainingExampleOutputParcel implements Parcelable {
         }
 
         /** Builds the instance. This builder should not be touched after calling this! */
-        public @android.annotation.NonNull TrainingExampleOutputParcel build() {
+        public @android.annotation.NonNull TrainingExamplesOutputParcel build() {
             checkNotUsed();
             mBuilderFieldsSet |= 0x4; // Mark builder used
 
@@ -205,9 +212,8 @@ public class TrainingExampleOutputParcel implements Parcelable {
             if ((mBuilderFieldsSet & 0x2) == 0) {
                 mResumptionTokens = null;
             }
-            TrainingExampleOutputParcel o = new TrainingExampleOutputParcel(
-                    mTrainingExamples,
-                    mResumptionTokens);
+            TrainingExamplesOutputParcel o =
+                    new TrainingExamplesOutputParcel(mTrainingExamples, mResumptionTokens);
             return o;
         }
 
@@ -220,10 +226,10 @@ public class TrainingExampleOutputParcel implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1695743444776L,
+            time = 1699394212826L,
             codegenVersion = "1.0.23",
             sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingExampleOutputParcel.java",
-            inputSignatures = " @android.annotation.Nullable com.android.ondevicepersonalization.internal.util.ByteArrayParceledListSlice mTrainingExamples\n @android.annotation.Nullable com.android.ondevicepersonalization.internal.util.ByteArrayParceledListSlice mResumptionTokens\nclass TrainingExampleOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genHiddenBuilder=true, genEqualsHashCode=true)")
+            inputSignatures = " @android.annotation.Nullable com.android.ondevicepersonalization.internal.util.ByteArrayParceledListSlice mTrainingExamples\n @android.annotation.Nullable com.android.ondevicepersonalization.internal.util.ByteArrayParceledListSlice mResumptionTokens\nclass TrainingExampleOutputParcel extends java.lang.Object implements [android.os.Parcelable]\n@com.android.ondevicepersonalization.internal.util.DataClass(genAidl=false, genHiddenBuilder=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java b/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
index eba12de1..890620e7 100644
--- a/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
+++ b/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
@@ -16,6 +16,9 @@
 
 package android.adservices.ondevicepersonalization;
 
+import static android.adservices.ondevicepersonalization.Constants.KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS;
+
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 
 import com.android.ondevicepersonalization.internal.util.AnnotationValidations;
@@ -23,37 +26,28 @@ import com.android.ondevicepersonalization.internal.util.DataClass;
 
 import java.time.Duration;
 
-/**
- * Training interval settings required for federated computation jobs.
- *
- * @hide
- */
+/** Training interval settings required for federated computation jobs. */
+@FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
 @DataClass(genBuilder = true, genHiddenConstDefs = true, genEqualsHashCode = true)
 public final class TrainingInterval {
-    /**
-     * The scheduling mode for a one-off task.
-     */
+    /** The scheduling mode for a one-off task. */
     public static final int SCHEDULING_MODE_ONE_TIME = 1;
 
-    /**
-     * The scheduling mode for a task that will be rescheduled after each run.
-     */
+    /** The scheduling mode for a task that will be rescheduled after each run. */
     public static final int SCHEDULING_MODE_RECURRENT = 2;
 
-
     /**
-     * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or
-     * {@link #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is
-     * {@link #SCHEDULING_MODE_ONE_TIME} if unspecified.
+     * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or {@link
+     * #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is {@link #SCHEDULING_MODE_ONE_TIME}
+     * if unspecified.
      */
     @SchedulingMode private int mSchedulingMode = SCHEDULING_MODE_ONE_TIME;
 
     /**
      * Sets the minimum time interval between two training runs.
      *
-     * <p>This field will only be used when the scheduling mode is
-     * {@link #SCHEDULING_MODE_RECURRENT}. Only positive values are accepted, zero or negative
-     * values will result in IllegalArgumentException.
+     * <p>This field will only be used when the scheduling mode is {@link
+     * #SCHEDULING_MODE_RECURRENT}. The value has be greater than zero.
      *
      * <p>Please also note this value is advisory, which does not guarantee the job will be run
      * immediately after the interval expired. Federated compute will still enforce a minimum
@@ -62,26 +56,23 @@ public final class TrainingInterval {
      */
     @NonNull private Duration mMinimumInterval = Duration.ZERO;
 
-
-
     // Code below generated by codegen v1.0.23.
     //
     // DO NOT MODIFY!
     // CHECKSTYLE:OFF Generated code
     //
     // To regenerate run:
-    // $ codegen $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
+    // $ codegen
+    // $ANDROID_BUILD_TOP/packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java
     //
     // To exclude the generated code from IntelliJ auto-formatting enable (one-time):
     //   Settings > Editor > Code Style > Formatter Control
-    //@formatter:off
-
+    // @formatter:off
 
     /** @hide */
-    @android.annotation.IntDef(prefix = "SCHEDULING_MODE_", value = {
-        SCHEDULING_MODE_ONE_TIME,
-        SCHEDULING_MODE_RECURRENT
-    })
+    @android.annotation.IntDef(
+            prefix = "SCHEDULING_MODE_",
+            value = {SCHEDULING_MODE_ONE_TIME, SCHEDULING_MODE_RECURRENT})
     @java.lang.annotation.Retention(java.lang.annotation.RetentionPolicy.SOURCE)
     @DataClass.Generated.Member
     public @interface SchedulingMode {}
@@ -91,38 +82,43 @@ public final class TrainingInterval {
     public static String schedulingModeToString(@SchedulingMode int value) {
         switch (value) {
             case SCHEDULING_MODE_ONE_TIME:
-                    return "SCHEDULING_MODE_ONE_TIME";
+                return "SCHEDULING_MODE_ONE_TIME";
             case SCHEDULING_MODE_RECURRENT:
-                    return "SCHEDULING_MODE_RECURRENT";
-            default: return Integer.toHexString(value);
+                return "SCHEDULING_MODE_RECURRENT";
+            default:
+                return Integer.toHexString(value);
         }
     }
 
     @DataClass.Generated.Member
     /* package-private */ TrainingInterval(
-            @SchedulingMode int schedulingMode,
-            @NonNull Duration minimumInterval) {
+            @SchedulingMode int schedulingMode, @NonNull Duration minimumInterval) {
         this.mSchedulingMode = schedulingMode;
 
         if (!(mSchedulingMode == SCHEDULING_MODE_ONE_TIME)
                 && !(mSchedulingMode == SCHEDULING_MODE_RECURRENT)) {
             throw new java.lang.IllegalArgumentException(
-                    "schedulingMode was " + mSchedulingMode + " but must be one of: "
-                            + "SCHEDULING_MODE_ONE_TIME(" + SCHEDULING_MODE_ONE_TIME + "), "
-                            + "SCHEDULING_MODE_RECURRENT(" + SCHEDULING_MODE_RECURRENT + ")");
+                    "schedulingMode was "
+                            + mSchedulingMode
+                            + " but must be one of: "
+                            + "SCHEDULING_MODE_ONE_TIME("
+                            + SCHEDULING_MODE_ONE_TIME
+                            + "), "
+                            + "SCHEDULING_MODE_RECURRENT("
+                            + SCHEDULING_MODE_RECURRENT
+                            + ")");
         }
 
         this.mMinimumInterval = minimumInterval;
-        AnnotationValidations.validate(
-                NonNull.class, null, mMinimumInterval);
+        AnnotationValidations.validate(NonNull.class, null, mMinimumInterval);
 
         // onConstructed(); // You can define this method to get a callback
     }
 
     /**
-     * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or
-     * {@link #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is
-     * {@link #SCHEDULING_MODE_ONE_TIME} if unspecified.
+     * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or {@link
+     * #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is {@link #SCHEDULING_MODE_ONE_TIME}
+     * if unspecified.
      */
     @DataClass.Generated.Member
     public @SchedulingMode int getSchedulingMode() {
@@ -132,9 +128,9 @@ public final class TrainingInterval {
     /**
      * Sets the minimum time interval between two training runs.
      *
-     * <p>This field will only be used when the scheduling mode is
-     * {@link #SCHEDULING_MODE_RECURRENT}. Only positive values are accepted, zero or negative
-     * values will result in IllegalArgumentException.
+     * <p>This field will only be used when the scheduling mode is {@link
+     * #SCHEDULING_MODE_RECURRENT}. Only positive values are accepted, zero or negative values will
+     * result in IllegalArgumentException.
      *
      * <p>Please also note this value is advisory, which does not guarantee the job will be run
      * immediately after the interval expired. Federated compute will still enforce a minimum
@@ -175,9 +171,8 @@ public final class TrainingInterval {
         return _hash;
     }
 
-    /**
-     * A builder for {@link TrainingInterval}
-     */
+    /** A builder for {@link TrainingInterval} */
+    @FlaggedApi(KEY_ENABLE_ONDEVICEPERSONALIZATION_APIS)
     @SuppressWarnings("WeakerAccess")
     @DataClass.Generated.Member
     public static final class Builder {
@@ -187,13 +182,12 @@ public final class TrainingInterval {
 
         private long mBuilderFieldsSet = 0L;
 
-        public Builder() {
-        }
+        public Builder() {}
 
         /**
-         * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or
-         * {@link #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is
-         * {@link #SCHEDULING_MODE_ONE_TIME} if unspecified.
+         * The scheduling mode for this task, either {@link #SCHEDULING_MODE_ONE_TIME} or {@link
+         * #SCHEDULING_MODE_RECURRENT}. The default scheduling mode is {@link
+         * #SCHEDULING_MODE_ONE_TIME} if unspecified.
          */
         @DataClass.Generated.Member
         public @NonNull Builder setSchedulingMode(@SchedulingMode int value) {
@@ -206,9 +200,9 @@ public final class TrainingInterval {
         /**
          * Sets the minimum time interval between two training runs.
          *
-         * <p>This field will only be used when the scheduling mode is
-         * {@link #SCHEDULING_MODE_RECURRENT}. Only positive values are accepted, zero or negative
-         * values will result in IllegalArgumentException.
+         * <p>This field will only be used when the scheduling mode is {@link
+         * #SCHEDULING_MODE_RECURRENT}. Only positive values are accepted, zero or negative values
+         * will result in IllegalArgumentException.
          *
          * <p>Please also note this value is advisory, which does not guarantee the job will be run
          * immediately after the interval expired. Federated compute will still enforce a minimum
@@ -234,9 +228,7 @@ public final class TrainingInterval {
             if ((mBuilderFieldsSet & 0x2) == 0) {
                 mMinimumInterval = Duration.ZERO;
             }
-            TrainingInterval o = new TrainingInterval(
-                    mSchedulingMode,
-                    mMinimumInterval);
+            TrainingInterval o = new TrainingInterval(mSchedulingMode, mMinimumInterval);
             return o;
         }
 
@@ -251,13 +243,14 @@ public final class TrainingInterval {
     @DataClass.Generated(
             time = 1697653739724L,
             codegenVersion = "1.0.23",
-            sourceFile = "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java",
-            inputSignatures = "public static final  int SCHEDULING_MODE_ONE_TIME\npublic static final  int SCHEDULING_MODE_RECURRENT\nprivate @android.adservices.ondevicepersonalization.TrainingInterval.SchedulingMode int mSchedulingMode\nprivate @android.annotation.NonNull java.time.Duration mMinimumInterval\nclass TrainingInterval extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genHiddenConstDefs=true, genEqualsHashCode=true)")
+            sourceFile =
+                    "packages/modules/OnDevicePersonalization/framework/java/android/adservices/ondevicepersonalization/TrainingInterval.java",
+            inputSignatures =
+                    "public static final  int SCHEDULING_MODE_ONE_TIME\npublic static final  int SCHEDULING_MODE_RECURRENT\nprivate @android.adservices.ondevicepersonalization.TrainingInterval.SchedulingMode int mSchedulingMode\nprivate @android.annotation.NonNull java.time.Duration mMinimumInterval\nclass TrainingInterval extends java.lang.Object implements []\n@com.android.ondevicepersonalization.internal.util.DataClass(genBuilder=true, genHiddenConstDefs=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
-
-    //@formatter:on
+    // @formatter:on
     // End of generated code
 
 }
diff --git a/src/com/android/ondevicepersonalization/services/display/DisplayHelper.java b/src/com/android/ondevicepersonalization/services/display/DisplayHelper.java
index b0aa954d..90a00847 100644
--- a/src/com/android/ondevicepersonalization/services/display/DisplayHelper.java
+++ b/src/com/android/ondevicepersonalization/services/display/DisplayHelper.java
@@ -16,7 +16,7 @@
 
 package com.android.ondevicepersonalization.services.display;
 
-import android.adservices.ondevicepersonalization.RenderOutput;
+import android.adservices.ondevicepersonalization.RenderOutputParcel;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.annotation.NonNull;
 import android.content.Context;
@@ -27,6 +27,7 @@ import android.os.PersistableBundle;
 import android.view.Display;
 import android.view.SurfaceControlViewHost;
 import android.view.SurfaceControlViewHost.SurfacePackage;
+import android.view.WindowManager;
 import android.webkit.WebSettings;
 import android.webkit.WebView;
 
@@ -58,9 +59,9 @@ public class DisplayHelper {
         mContext = context;
     }
 
-    /** Generates an HTML string from the template data in RenderOutput. */
+    /** Generates an HTML string from the template data in RenderOutputParcel. */
     @NonNull public String generateHtml(
-            @NonNull RenderOutput renderContentResult,
+            @NonNull RenderOutputParcel renderContentResult,
             @NonNull String servicePackageName) {
         // If htmlContent is provided, do not render the template.
         String htmlContent = renderContentResult.getContent();
@@ -134,7 +135,12 @@ public class DisplayHelper {
             @NonNull SettableFuture<SurfacePackage> resultFuture) {
         try {
             sLogger.d(TAG + ": createWebView() started");
-            WebView webView = new WebView(mContext);
+            Display display = mContext.getSystemService(DisplayManager.class).getDisplay(displayId);
+            Context displayContext = mContext.createDisplayContext(display);
+            Context windowContext = displayContext.createWindowContext(
+                    WindowManager.LayoutParams.TYPE_APPLICATION_PANEL, null);
+
+            WebView webView = new WebView(windowContext);
             webView.setWebViewClient(
                     new OdpWebViewClient(mContext, servicePackageName, queryId, logRecord));
             WebSettings webViewSettings = webView.getSettings();
@@ -143,8 +149,8 @@ public class DisplayHelper {
             webViewSettings.setAllowContentAccess(false);
             webView.loadData(html, "text/html; charset=utf-8", "UTF-8");
 
-            Display display = mContext.getSystemService(DisplayManager.class).getDisplay(displayId);
-            SurfaceControlViewHost host = new SurfaceControlViewHost(mContext, display, hostToken);
+            SurfaceControlViewHost host = new SurfaceControlViewHost(
+                    windowContext, display, hostToken);
             host.setView(webView, width, height);
             SurfacePackage surfacePackage = host.getSurfacePackage();
             sLogger.d(TAG + ": createWebView success: " + surfacePackage);
diff --git a/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java b/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
index 0be485ef..d74f0c71 100644
--- a/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
+++ b/src/com/android/ondevicepersonalization/services/display/OdpWebViewClient.java
@@ -17,9 +17,9 @@
 package com.android.ondevicepersonalization.services.display;
 
 import android.adservices.ondevicepersonalization.Constants;
-import android.adservices.ondevicepersonalization.EventInput;
+import android.adservices.ondevicepersonalization.EventInputParcel;
 import android.adservices.ondevicepersonalization.EventLogRecord;
-import android.adservices.ondevicepersonalization.EventOutput;
+import android.adservices.ondevicepersonalization.EventOutputParcel;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.adservices.ondevicepersonalization.UserData;
 import android.annotation.NonNull;
@@ -181,7 +181,7 @@ class OdpWebViewClient extends WebViewClient {
         return true;
     }
 
-    private ListenableFuture<EventOutput> executeEventHandler(
+    private ListenableFuture<EventOutputParcel> executeEventHandler(
             IsolatedServiceInfo isolatedServiceInfo,
             EventUrlPayload payload) {
         try {
@@ -192,7 +192,7 @@ class OdpWebViewClient extends WebViewClient {
                     /* includeEventData */ true);
             serviceParams.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, binder);
             // TODO(b/259950177): Add Query row to input.
-            EventInput input = new EventInput.Builder()
+            EventInputParcel input = new EventInputParcel.Builder()
                     .setParameters(payload.getEventParams())
                     .setRequestLogRecord(mLogRecord)
                     .build();
@@ -213,7 +213,7 @@ class OdpWebViewClient extends WebViewClient {
                                         result, isolatedServiceInfo.getStartTimeMillis(),
                                         Constants.STATUS_SUCCESS);
                                 return result.getParcelable(
-                                        Constants.EXTRA_RESULT, EventOutput.class);
+                                        Constants.EXTRA_RESULT, EventOutputParcel.class);
                             },
                             mInjector.getExecutor())
                     .catchingAsync(
@@ -233,25 +233,25 @@ class OdpWebViewClient extends WebViewClient {
 
     }
 
-    ListenableFuture<EventOutput> getEventOutput(
+    ListenableFuture<EventOutputParcel> getEventOutputParcel(
             ListenableFuture<IsolatedServiceInfo> loadFuture,
             EventUrlPayload payload) {
         try {
-            sLogger.d(TAG + ": getEventOutput(): Starting isolated process.");
+            sLogger.d(TAG + ": getEventOutputParcel(): Starting isolated process.");
             return FluentFuture.from(loadFuture)
                 .transformAsync(
                         result -> executeEventHandler(result, payload),
                         mInjector.getExecutor());
 
         } catch (Exception e) {
-            sLogger.e(TAG + ": getEventOutput() failed", e);
+            sLogger.e(TAG + ": getEventOutputParcel() failed", e);
             return Futures.immediateFailedFuture(e);
         }
     }
 
-    private ListenableFuture<Void> writeEvent(EventOutput result) {
+    private ListenableFuture<Void> writeEvent(EventOutputParcel result) {
         try {
-            sLogger.d(TAG + ": writeEvent() called. EventOutput: " + result.toString());
+            sLogger.d(TAG + ": writeEvent() called. EventOutputParcel: " + result.toString());
             if (result == null || result.getEventLogRecord() == null) {
                 return Futures.immediateFuture(null);
             }
@@ -293,7 +293,7 @@ class OdpWebViewClient extends WebViewClient {
                     mInjector.getProcessRunner().loadIsolatedService(
                         TASK_NAME, mServicePackageName);
 
-            var doneFuture = FluentFuture.from(getEventOutput(loadFuture, eventUrlPayload))
+            var doneFuture = FluentFuture.from(getEventOutputParcel(loadFuture, eventUrlPayload))
                     .transformAsync(
                         result -> writeEvent(result),
                         mInjector.getExecutor())
diff --git a/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java b/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
index 9156807b..b05527e4 100644
--- a/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
+++ b/src/com/android/ondevicepersonalization/services/download/OnDevicePersonalizationDataProcessingAsyncCallable.java
@@ -17,7 +17,7 @@
 package com.android.ondevicepersonalization.services.download;
 
 import android.adservices.ondevicepersonalization.Constants;
-import android.adservices.ondevicepersonalization.DownloadCompletedOutput;
+import android.adservices.ondevicepersonalization.DownloadCompletedOutputParcel;
 import android.adservices.ondevicepersonalization.DownloadInputParcel;
 import android.adservices.ondevicepersonalization.UserData;
 import android.content.Context;
@@ -236,8 +236,8 @@ public class OnDevicePersonalizationDataProcessingAsyncCallable implements Async
             Map<String, VendorData> vendorDataMap) {
         sLogger.d(TAG + ": Plugin filter code completed successfully");
         List<VendorData> filteredList = new ArrayList<>();
-        DownloadCompletedOutput downloadResult = pluginResult.getParcelable(
-                Constants.EXTRA_RESULT, DownloadCompletedOutput.class);
+        DownloadCompletedOutputParcel downloadResult = pluginResult.getParcelable(
+                Constants.EXTRA_RESULT, DownloadCompletedOutputParcel.class);
         List<String> retainedKeys = downloadResult.getRetainedKeys();
         if (retainedKeys == null) {
             // TODO(b/270710021): Determine how to correctly handle null retainedKeys.
diff --git a/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java b/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
index b0ae9bd4..5d63c651 100644
--- a/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
+++ b/src/com/android/ondevicepersonalization/services/federatedcompute/OdpExampleStoreService.java
@@ -17,8 +17,8 @@
 package com.android.ondevicepersonalization.services.federatedcompute;
 
 import android.adservices.ondevicepersonalization.Constants;
-import android.adservices.ondevicepersonalization.TrainingExampleInput;
-import android.adservices.ondevicepersonalization.TrainingExampleOutputParcel;
+import android.adservices.ondevicepersonalization.TrainingExamplesInputParcel;
+import android.adservices.ondevicepersonalization.TrainingExamplesOutputParcel;
 import android.adservices.ondevicepersonalization.UserData;
 import android.annotation.NonNull;
 import android.content.Context;
@@ -146,8 +146,8 @@ public final class OdpExampleStoreService extends ExampleStoreService {
                 resumptionToken = eventState.getToken();
             }
 
-            TrainingExampleInput input =
-                    new TrainingExampleInput.Builder()
+            TrainingExamplesInputParcel input =
+                    new TrainingExamplesInputParcel.Builder()
                             .setResumptionToken(resumptionToken)
                             .setPopulationName(populationName)
                             .setTaskName(taskName)
@@ -155,16 +155,16 @@ public final class OdpExampleStoreService extends ExampleStoreService {
 
             ListenableFuture<IsolatedServiceInfo> loadFuture =
                     mInjector.getProcessRunner().loadIsolatedService(TASK_NAME, packageName);
-            ListenableFuture<TrainingExampleOutputParcel> resultFuture =
+            ListenableFuture<TrainingExamplesOutputParcel> resultFuture =
                     FluentFuture.from(loadFuture)
                             .transformAsync(
-                                    result -> executeOnTrainingExample(result, input, packageName),
+                                    result -> executeOnTrainingExamples(result, input, packageName),
                                     OnDevicePersonalizationExecutors.getBackgroundExecutor())
                             .transform(
                                     result -> {
                                         return result.getParcelable(
                                                 Constants.EXTRA_RESULT,
-                                                TrainingExampleOutputParcel.class);
+                                                TrainingExamplesOutputParcel.class);
                                     },
                                     OnDevicePersonalizationExecutors.getBackgroundExecutor())
                             .withTimeout(
@@ -174,14 +174,14 @@ public final class OdpExampleStoreService extends ExampleStoreService {
 
             Futures.addCallback(
                     resultFuture,
-                    new FutureCallback<TrainingExampleOutputParcel>() {
+                    new FutureCallback<TrainingExamplesOutputParcel>() {
                         @Override
                         public void onSuccess(
-                                TrainingExampleOutputParcel trainingExampleOutputParcel) {
+                                TrainingExamplesOutputParcel trainingExamplesOutputParcel) {
                             ByteArrayParceledListSlice trainingExamplesListSlice =
-                                    trainingExampleOutputParcel.getTrainingExamples();
+                                    trainingExamplesOutputParcel.getTrainingExamples();
                             ByteArrayParceledListSlice resumptionTokensListSlice =
-                                    trainingExampleOutputParcel.getResumptionTokens();
+                                    trainingExamplesOutputParcel.getResumptionTokens();
                             if (trainingExamplesListSlice == null
                                     || resumptionTokensListSlice == null) {
                                 callback.onStartQuerySuccess(
@@ -219,11 +219,11 @@ public final class OdpExampleStoreService extends ExampleStoreService {
         }
     }
 
-    private ListenableFuture<Bundle> executeOnTrainingExample(
+    private ListenableFuture<Bundle> executeOnTrainingExamples(
             IsolatedServiceInfo isolatedServiceInfo,
-            TrainingExampleInput exampleInput,
+            TrainingExamplesInputParcel exampleInput,
             String packageName) {
-        sLogger.d(TAG + ": executeOnTrainingExample() started.");
+        sLogger.d(TAG + ": executeOnTrainingExamples() started.");
         Bundle serviceParams = new Bundle();
         serviceParams.putParcelable(Constants.EXTRA_INPUT, exampleInput);
         DataAccessServiceImpl binder =
diff --git a/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java b/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
index b5f6288a..88d94214 100644
--- a/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
+++ b/src/com/android/ondevicepersonalization/services/request/AppRequestFlow.java
@@ -18,8 +18,8 @@ package com.android.ondevicepersonalization.services.request;
 
 import android.adservices.ondevicepersonalization.Constants;
 import android.adservices.ondevicepersonalization.EventLogRecord;
-import android.adservices.ondevicepersonalization.ExecuteInput;
-import android.adservices.ondevicepersonalization.ExecuteOutput;
+import android.adservices.ondevicepersonalization.ExecuteInputParcel;
+import android.adservices.ondevicepersonalization.ExecuteOutputParcel;
 import android.adservices.ondevicepersonalization.RenderingConfig;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.adservices.ondevicepersonalization.UserData;
@@ -178,7 +178,7 @@ public class AppRequestFlow {
             ListenableFuture<IsolatedServiceInfo> loadFuture =
                     mInjector.getProcessRunner().loadIsolatedService(
                         TASK_NAME, mService.getPackageName());
-            ListenableFuture<ExecuteOutput> resultFuture = FluentFuture.from(loadFuture)
+            ListenableFuture<ExecuteOutputParcel> resultFuture = FluentFuture.from(loadFuture)
                     .transformAsync(
                             result -> executeAppRequest(result),
                             mInjector.getExecutor()
@@ -186,7 +186,7 @@ public class AppRequestFlow {
                     .transform(
                             result -> {
                                 return result.getParcelable(
-                                        Constants.EXTRA_RESULT, ExecuteOutput.class);
+                                        Constants.EXTRA_RESULT, ExecuteOutputParcel.class);
                             },
                             mInjector.getExecutor()
                     );
@@ -239,8 +239,8 @@ public class AppRequestFlow {
             IsolatedServiceInfo isolatedServiceInfo) {
         sLogger.d(TAG + ": executeAppRequest() started.");
         Bundle serviceParams = new Bundle();
-        ExecuteInput input =
-                new ExecuteInput.Builder()
+        ExecuteInputParcel input =
+                new ExecuteInputParcel.Builder()
                         .setAppPackageName(mCallingPackageName)
                         .setAppParams(mParams)
                         .build();
@@ -279,7 +279,7 @@ public class AppRequestFlow {
                 );
     }
 
-    private ListenableFuture<Long> logQuery(ExecuteOutput result) {
+    private ListenableFuture<Long> logQuery(ExecuteOutputParcel result) {
         sLogger.d(TAG + ": logQuery() started.");
         EventsDao eventsDao = EventsDao.getInstance(mContext);
         // Insert query
@@ -335,11 +335,11 @@ public class AppRequestFlow {
     }
 
     private ListenableFuture<List<String>> createTokens(
-            ListenableFuture<ExecuteOutput> resultFuture,
+            ListenableFuture<ExecuteOutputParcel> resultFuture,
             ListenableFuture<Long> queryIdFuture) {
         try {
             sLogger.d(TAG + ": createTokens() started.");
-            ExecuteOutput result = Futures.getDone(resultFuture);
+            ExecuteOutputParcel result = Futures.getDone(resultFuture);
             long queryId = Futures.getDone(queryIdFuture);
             List<RenderingConfig> renderingConfigs = result.getRenderingConfigs();
             Objects.requireNonNull(renderingConfigs);
diff --git a/src/com/android/ondevicepersonalization/services/request/RenderFlow.java b/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
index 6ba342ed..d88e67a5 100644
--- a/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
+++ b/src/com/android/ondevicepersonalization/services/request/RenderFlow.java
@@ -17,8 +17,8 @@
 package com.android.ondevicepersonalization.services.request;
 
 import android.adservices.ondevicepersonalization.Constants;
-import android.adservices.ondevicepersonalization.RenderInput;
-import android.adservices.ondevicepersonalization.RenderOutput;
+import android.adservices.ondevicepersonalization.RenderInputParcel;
+import android.adservices.ondevicepersonalization.RenderOutputParcel;
 import android.adservices.ondevicepersonalization.RenderingConfig;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.adservices.ondevicepersonalization.aidl.IRequestSurfacePackageCallback;
@@ -229,7 +229,7 @@ public class RenderFlow {
                             mInjector.getExecutor())
                     .transform(result -> {
                         return result.getParcelable(
-                                Constants.EXTRA_RESULT, RenderOutput.class);
+                                Constants.EXTRA_RESULT, RenderOutputParcel.class);
                     }, mInjector.getExecutor())
                     .transform(
                             result -> mDisplayHelper.generateHtml(result, mServicePackageName),
@@ -255,8 +255,8 @@ public class RenderFlow {
             RenderingConfig renderingConfig) {
         sLogger.d(TAG + "executeRenderContentRequest() started.");
         Bundle serviceParams = new Bundle();
-        RenderInput input =
-                new RenderInput.Builder()
+        RenderInputParcel input =
+                new RenderInputParcel.Builder()
                     .setHeight(mHeight)
                     .setWidth(mWidth)
                     .setRenderingConfigIndex(slotIndex)
diff --git a/tests/endtoendtests/OdpClient/AndroidManifest.xml b/tests/endtoendtests/OdpClient/AndroidManifest.xml
index 71025450..1af14085 100644
--- a/tests/endtoendtests/OdpClient/AndroidManifest.xml
+++ b/tests/endtoendtests/OdpClient/AndroidManifest.xml
@@ -17,13 +17,14 @@
 
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
           package="com.example.odpclient"
-          android:versionName="1.0.2" >
+          android:versionName="1.0.3" >
     <uses-permission android:name="android.permission.ondevicepersonalization.MODIFY_ONDEVICEPERSONALIZATION_STATE"/>
     <application
         android:label="@string/title_activity_main">
 
         <activity
             android:name="com.example.odpclient.MainActivity"
+            android:configChanges="orientation|screenSize|screenLayout|keyboardHidden"
             android:exported="true">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
diff --git a/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java b/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
index 2fe34124..7e445d89 100644
--- a/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
+++ b/tests/endtoendtests/OdpClient/src/com/example/odpclient/MainActivity.java
@@ -22,6 +22,7 @@ import android.adservices.ondevicepersonalization.SurfacePackageToken;
 import android.app.Activity;
 import android.content.ComponentName;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
@@ -29,6 +30,7 @@ import android.os.OutcomeReceiver;
 import android.os.PersistableBundle;
 import android.util.Log;
 import android.view.SurfaceControlViewHost.SurfacePackage;
+import android.view.SurfaceHolder;
 import android.view.SurfaceView;
 import android.view.View;
 import android.widget.Button;
@@ -45,7 +47,6 @@ import java.util.concurrent.atomic.AtomicReference;
 
 public class MainActivity extends Activity {
     private static final String TAG = "OdpClient";
-    private OnDevicePersonalizationManager mOdpManager = null;
     private OnDevicePersonalizationConfigManager mOdpConfigManager = null;
 
     private EditText mTextBox;
@@ -59,20 +60,32 @@ public class MainActivity extends Activity {
     private Context mContext;
     private static Executor sCallbackExecutor = Executors.newSingleThreadExecutor();
 
+    class SurfaceCallback implements SurfaceHolder.Callback {
+        @Override public void surfaceCreated(SurfaceHolder holder) {
+            Log.d(TAG, "surfaceCreated");
+        }
+        @Override public void surfaceDestroyed(SurfaceHolder holder) {
+            Log.d(TAG, "surfaceDestroyed");
+        }
+        @Override public void surfaceChanged(
+                SurfaceHolder holder, int format, int width, int height) {
+            Log.d(TAG, "surfaceChanged");
+        }
+    }
+
     @Override
     public void onCreate(Bundle savedInstanceState) {
+        Log.d(TAG, "onCreate");
         super.onCreate(savedInstanceState);
         setContentView(R.layout.activity_main);
         mContext = getApplicationContext();
-        if (mOdpManager == null) {
-            mOdpManager = mContext.getSystemService(OnDevicePersonalizationManager.class);
-        }
         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);
@@ -99,18 +112,19 @@ public class MainActivity extends Activity {
         mReportConversionButton.setOnClickListener(v -> reportConversion());
     }
 
+    private OnDevicePersonalizationManager getOdpManager() {
+        return mContext.getSystemService(OnDevicePersonalizationManager.class);
+    }
+
     private void makeRequest() {
         try {
-            if (mOdpManager == null) {
-                makeToast("OnDevicePersonalizationManager is null");
-                return;
-            }
+            var odpManager = getOdpManager();
             CountDownLatch latch = new CountDownLatch(1);
             Log.i(TAG, "Starting execute()");
             AtomicReference<SurfacePackageToken> slotResultHandle = new AtomicReference<>();
             PersistableBundle appParams = new PersistableBundle();
             appParams.putString("keyword", mTextBox.getText().toString());
-            mOdpManager.execute(
+            odpManager.execute(
                     ComponentName.createRelative(
                         "com.example.odpsamplenetwork",
                         "com.example.odpsamplenetwork.SampleService"),
@@ -136,7 +150,7 @@ public class MainActivity extends Activity {
                     });
             latch.await();
             Log.d(TAG, "wait success");
-            mOdpManager.requestSurfacePackage(
+            odpManager.requestSurfacePackage(
                     slotResultHandle.get(),
                     mRenderedView.getHostToken(),
                     getDisplay().getDisplayId(),
@@ -176,15 +190,12 @@ public class MainActivity extends Activity {
 
     private void scheduleTraining() {
         try {
-            if (mOdpManager == null) {
-                makeToast("OnDevicePersonalizationManager is null");
-                return;
-            }
+            var odpManager = getOdpManager();
             CountDownLatch latch = new CountDownLatch(1);
             Log.i(TAG, "Starting execute()");
             PersistableBundle appParams = new PersistableBundle();
             appParams.putString("schedule_training", mScheduleTrainingTextBox.getText().toString());
-            mOdpManager.execute(
+            odpManager.execute(
                     ComponentName.createRelative(
                             "com.example.odpsamplenetwork",
                             "com.example.odpsamplenetwork.SampleService"),
@@ -211,15 +222,12 @@ public class MainActivity extends Activity {
 
     private void reportConversion() {
         try {
-            if (mOdpManager == null) {
-                makeToast("OnDevicePersonalizationManager is null");
-                return;
-            }
+            var odpManager = getOdpManager();
             CountDownLatch latch = new CountDownLatch(1);
             Log.i(TAG, "Starting execute()");
             PersistableBundle appParams = new PersistableBundle();
             appParams.putString("conversion_ad_id", mReportConversionTextBox.getText().toString());
-            mOdpManager.execute(
+            odpManager.execute(
                     ComponentName.createRelative(
                             "com.example.odpsamplenetwork",
                             "com.example.odpsamplenetwork.SampleService"),
@@ -268,4 +276,39 @@ public class MainActivity extends Activity {
         Log.i(TAG, message);
         runOnUiThread(() -> Toast.makeText(MainActivity.this, message, Toast.LENGTH_LONG).show());
     }
+
+    @Override
+    public void onPause() {
+        Log.d(TAG, "onPause");
+        super.onPause();
+    }
+    @Override
+    public void onSaveInstanceState(Bundle outState) {
+        Log.d(TAG, "onSaveInstanceState");
+        super.onSaveInstanceState(outState);
+    }
+
+    @Override
+    public void onDestroy() {
+        Log.d(TAG, "onDestroy");
+        super.onDestroy();
+    }
+
+    @Override
+    public void onRestoreInstanceState(Bundle savedInstanceState) {
+        Log.d(TAG, "onRestoreInstanceState");
+        super.onRestoreInstanceState(savedInstanceState);
+    }
+
+    @Override
+    public void onResume() {
+        Log.d(TAG, "onResume");
+        super.onResume();
+    }
+
+    @Override
+    public void onConfigurationChanged(Configuration newConfig) {
+        Log.d(TAG, "onConfigurationChanged");
+        super.onConfigurationChanged(newConfig);
+    }
 }
diff --git a/tests/endtoendtests/OdpSampleNetwork/AndroidManifest.xml b/tests/endtoendtests/OdpSampleNetwork/AndroidManifest.xml
index 3dce1e63..0d98ec4f 100644
--- a/tests/endtoendtests/OdpSampleNetwork/AndroidManifest.xml
+++ b/tests/endtoendtests/OdpSampleNetwork/AndroidManifest.xml
@@ -17,7 +17,7 @@
 
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
           package="com.example.odpsamplenetwork"
-          android:versionName="1.0.2" >
+          android:versionName="1.0.3" >
     <application android:label="OdpSampleNetwork"
                  android:debuggable="true">
         <property android:name="android.ondevicepersonalization.ON_DEVICE_PERSONALIZATION_CONFIG"
diff --git a/tests/endtoendtests/OdpSampleNetwork/res/raw/test_data1.json b/tests/endtoendtests/OdpSampleNetwork/res/raw/test_data1.json
index 6cd2a147..c704637b 100644
--- a/tests/endtoendtests/OdpSampleNetwork/res/raw/test_data1.json
+++ b/tests/endtoendtests/OdpSampleNetwork/res/raw/test_data1.json
@@ -6,6 +6,106 @@
     },
     { "key": "template1",
       "data": "<img src=\"$impressionUrl\">\n<a href=\"$clickUrl\">${adText}!</a>"
-    }
+    },
+    {"key": "example1", "data": "0,1,1,5,0,1382,4,15,2,181,1,2,,2,68fd1e64,80e26c9b,fb936136,7b4723c4,25c83c98,7e0ccccf,de7995b8,1f89b562,a73ee510,a8cd5504,b2cb9c98,37c9c164,2824a5f6,1adce6ef,8ba8b39a,891b62e7,e5ba7672,f54016b9,21ddcdc9,b1252a9d,07b5194c,,3a171ecb,c5c50484,e8b83407,9727dd16,"},
+    {"key": "example2", "data": "0,2,0,44,1,102,8,2,2,4,1,1,,4,68fd1e64,f0cf0024,6f67f7e5,41274cd7,25c83c98,fe6b92e5,922afcc0,0b153874,a73ee510,2b53e5fb,4f1b46f3,623049e6,d7020589,b28479f6,e6c5b5cd,c92f3b61,07c540c4,b04e4670,21ddcdc9,5840adea,60f6221e,,3a171ecb,43f13e8b,e8b83407,731c3655,"},
+    {"key": "example3", "data": "0,2,0,1,14,767,89,4,2,245,1,3,3,45,287e684f,0a519c5c,02cf9876,c18be181,25c83c98,7e0ccccf,c78204a1,0b153874,a73ee510,3b08e48b,5f5e6091,8fe001f4,aa655a2f,07d13a8f,6dc710ed,36103458,8efede7f,3412118d,,,e587c466,ad3062eb,3a171ecb,3b183c5c,,,"},
+    {"key": "example4", "data": "0,,893,,,4392,,0,0,0,,0,,,68fd1e64,2c16a946,a9a87e68,2e17d6f6,25c83c98,fe6b92e5,2e8a689b,0b153874,a73ee510,efea433b,e51ddf94,a30567ca,3516f6e6,07d13a8f,18231224,52b8680f,1e88c74f,74ef3502,,,6b3a5ca6,,3a171ecb,9117a34a,,,"},
+    {"key": "example5", "data": "0,3,-1,,0,2,0,3,0,0,1,1,,0,8cf07265,ae46a29d,c81688bb,f922efad,25c83c98,13718bbd,ad9fa255,0b153874,a73ee510,5282c137,e5d8af57,66a76a26,f06c53ac,1adce6ef,8ff4b403,01adbab4,1e88c74f,26b3c7a7,,,21c9516a,,32c7478e,b34f3128,,,"},
+    {"key": "example6", "data": "0,,-1,,,12824,,0,0,6,,0,,,05db9164,6c9c9cf3,2730ec9c,5400db8b,43b19349,6f6d9be8,53b5f978,0b153874,a73ee510,3b08e48b,91e8fc27,be45b877,9ff13f22,07d13a8f,06969a20,9bc7fff5,776ce399,92555263,,,242bb710,8ec974f4,be7c41b4,72c78f11,,,"},
+    {"key": "example7", "data": "0,,1,2,,3168,,0,1,2,,0,,,439a44a4,ad4527a2,c02372d0,d34ebbaa,43b19349,fe6b92e5,4bc6ffea,0b153874,a73ee510,3b08e48b,a4609aab,14d63538,772a00d7,07d13a8f,f9d1382e,b00d3dc9,776ce399,cdfa8259,,,20062612,,93bad2c0,1b256e61,,,"},
+    {"key": "example8", "data": "1,1,4,2,0,0,0,1,0,0,1,1,,0,68fd1e64,2c16a946,503b9dbc,e4dbea90,f3474129,13718bbd,38eb9cf4,1f89b562,a73ee510,547c0ffe,bc8c9f21,60ab2f07,46f42a63,07d13a8f,18231224,e6b6bdc7,e5ba7672,74ef3502,,,5316a17f,,32c7478e,9117a34a,,,"},
+    {"key": "example9", "data": "0,,44,4,8,19010,249,28,31,141,,1,,8,05db9164,d833535f,d032c263,c18be181,25c83c98,7e0ccccf,d5b6acf2,0b153874,a73ee510,2acdcf4e,086ac2d2,dfbb09fb,41a6ae00,b28479f6,e2502ec9,84898b2a,e5ba7672,42a2edb9,,,0014c32a,,32c7478e,3b183c5c,,,"},
+    {"key": "example10", "data": "0,,35,,1,33737,21,1,2,3,,1,,1,05db9164,510b40a5,d03e7c24,eb1fd928,25c83c98,,52283d1c,0b153874,a73ee510,015ac893,e51ddf94,951fe4a9,3516f6e6,07d13a8f,2ae4121c,8ec71479,d4bb7bd8,70d0f5f9,,,0e63fca0,,32c7478e,0e8fe315,,,"},
+    {"key": "example11", "data": "0,,2,632,0,56770,,0,5,65,,0,,2,05db9164,0468d672,7ae80d0f,80d8555a,25c83c98,7e0ccccf,04277bf9,0b153874,7cc72ec2,3b08e48b,7e2c5c15,cfc86806,91a1b611,b28479f6,58251aab,146a70fd,776ce399,0b331314,21ddcdc9,5840adea,cbec39db,,3a171ecb,cedad179,ea9a246c,9a556cfc,"},
+    {"key": "example12", "data": "0,0,6,6,6,421,109,1,7,107,0,1,,6,05db9164,9b5fd12f,,,4cf72387,,111121f4,0b153874,a73ee510,3b08e48b,ac9c2e8f,,6e2d6a15,07d13a8f,796a1a2e,,d4bb7bd8,8aaa5b67,,,,,32c7478e,,,,"},
+    {"key": "example13", "data": "1,0,-1,,,1465,0,17,0,4,0,4,,,241546e0,38a947a1,fa673455,6a14f9b9,25c83c98,fe6b92e5,1c86e0eb,1f89b562,a73ee510,e7ba2569,755e4a50,208d9687,5978055e,07d13a8f,5182f694,f8b34416,e5ba7672,e5f8f18f,,,f3ddd519,,32c7478e,b34f3128,,,"},
+    {"key": "example14", "data": "1,,2,11,5,10262,34,2,4,5,,1,,5,be589b51,287130e0,cd7a7a22,fb7334df,25c83c98,,6cdb3998,361384ce,a73ee510,3ff10fb2,5874c9c9,976cbd4c,740c210d,1adce6ef,310d155b,07eb8110,07c540c4,891589e7,18259a83,a458ea53,a0ab60ca,,32c7478e,a052b1ed,9b3e8820,8967c0d2,"},
+    {"key": "example15", "data": "0,0,51,84,4,3633,26,1,4,8,0,1,,4,5a9ed9b0,80e26c9b,97144401,5dbf0cc5,0942e0a7,13718bbd,9ce6136d,0b153874,a73ee510,2106e595,b5bb9d63,04f55317,ab04d8fe,1adce6ef,0ad47a49,2bd32e5c,3486227d,12195b22,21ddcdc9,b1252a9d,fa131867,,dbb486d7,8ecc176a,e8b83407,c43c3f58,"},
+    {"key": "example16", "data": "0,,2,1,18,20255,,0,1,1306,,0,,20,05db9164,bc6e3dc1,67799c69,d00d0f35,4cf72387,7e0ccccf,ca4fd8f8,64523cfa,a73ee510,3b08e48b,a0060bca,b9f28c33,22d23aac,5aebfb83,d702713a,0f655650,776ce399,3a2028fd,,,b426bc93,,3a171ecb,2e0a0035,,,"},
+    {"key": "example17", "data": "1,1,987,,2,105,2,1,2,2,1,1,,2,68fd1e64,38d50e09,da603082,431a5096,43b19349,7e0ccccf,3f35b640,0b153874,a73ee510,3b08e48b,3d5fb018,6aaab577,94172618,07d13a8f,ee569ce2,2f03ef40,d4bb7bd8,582152eb,21ddcdc9,b1252a9d,3b203ca1,,32c7478e,b21dc903,001f3601,aa5f0a15,"},
+    {"key": "example18", "data": "0,0,1,,0,16597,557,3,5,123,0,1,,1,8cf07265,7cd19acc,77f2f2e5,d16679b9,4cf72387,fbad5c96,8fb24933,0b153874,a73ee510,0095a535,3617b5f5,9f32b866,428332cf,b28479f6,83ebd498,31ca40b6,e5ba7672,d0e5eb07,,,dfcfc3fa,ad3062eb,32c7478e,aee52b6f,,,"},
+    {"key": "example19", "data": "0,0,24,4,2,2056,12,6,10,83,0,1,,2,05db9164,f0cf0024,08b45d8b,cbb5af1b,384874ce,fbad5c96,81bb0302,37e4aa92,a73ee510,175d6c71,b7094596,1c547463,1f9d2c38,1adce6ef,55dc357b,0ca69655,e5ba7672,b04e4670,21ddcdc9,b1252a9d,f3caefdd,,32c7478e,4c8e5aef,ea9a246c,9593bba9,"},
+    {"key": "example20", "data": "0,7,102,,3,780,15,7,15,15,1,1,,3,3c9d8785,b0660259,3a960356,15c92ddb,4cf72387,13718bbd,00c46cd1,0b153874,a73ee510,62cfc6bd,8cffe207,656e5413,ff5626de,ad1cc976,27b1230c,fa8d05aa,e5ba7672,5edd90de,,,e12ce348,,c3dc6cef,49045073,,,"},
+    {"key": "example21", "data": "1,,47,,0,6399,38,19,10,143,,10,,6,1464facd,38a947a1,223b0e16,ca55061c,25c83c98,7e0ccccf,6933dec1,5b392875,a73ee510,3b08e48b,860c302b,156f99ef,30735474,1adce6ef,0e78291e,5fbf4a84,e5ba7672,1999bae9,,,deb9605d,,32c7478e,e448275f,,,"},
+    {"key": "example22", "data": "0,0,1,80,0,1848,287,1,4,46,0,1,,4,05db9164,09e68b86,13b87f72,13a91973,25c83c98,7e0ccccf,cc5ed2f1,0b153874,a73ee510,3b08e48b,081c279a,d25f00b6,9f16a973,07d13a8f,36721ddc,1746d357,d4bb7bd8,5aed7436,a153cea2,a458ea53,dd37e0d1,,32c7478e,c70a58f2,e8b83407,af7ece63,"},
+    {"key": "example23", "data": "0,,0,14,6,7132,171,2,2,6,,1,,6,05db9164,38a947a1,e88a1d4c,8eb9aec7,25c83c98,fbad5c96,3fd38f3b,5b392875,a73ee510,5162b19c,7c430b79,4ac05ba7,7f0d7407,b28479f6,d1128331,ce881087,07c540c4,5d93f8ab,,,57d0811b,,3a171ecb,1793a828,,,"},
+    {"key": "example24", "data": "0,,9,9,17,11774,,0,23,128,,0,,17,05db9164,08d6d899,cf59444f,60d5f5a7,25c83c98,7e0ccccf,38850d41,0b153874,a73ee510,6e7947ce,49aeb6a9,1d00cbc4,8f7e5dc7,07d13a8f,41f10449,b93ac0ad,1e88c74f,698d1c68,,,bf8efd4c,,c7dc6720,f96a556f,,,"},
+    {"key": "example25", "data": "0,0,1,2,,6190,84,1,27,71,0,1,,,5a9ed9b0,3df44d94,d032c263,c18be181,25c83c98,7e0ccccf,a0845add,0b153874,a73ee510,967857d1,e469acef,dfbb09fb,849a0a56,07d13a8f,72d05a1c,84898b2a,d4bb7bd8,e7648a8f,,,0014c32a,c9d4222a,3a171ecb,3b183c5c,,,"},
+    {"key": "example26", "data": "0,,4,16,,5925,2,2,0,0,,1,,,5a9ed9b0,09e68b86,64094ddd,b0a4d1e3,25c83c98,,b87f4a4a,0b153874,a73ee510,2124a520,319687c9,b51dc799,62036f49,64c94865,91126f30,1a00d73c,07c540c4,5aed7436,aa0cf899,a458ea53,c30dce78,,32c7478e,3fdb382b,e8b83407,49d68486,"},
+    {"key": "example27", "data": "1,0,1,20,16,1548,93,42,32,912,0,15,1,16,8cf07265,942f9a8d,a8e40bcf,0365276a,25c83c98,7e0ccccf,3f4ec687,1f89b562,a73ee510,726f00fd,c4adf918,27c604a6,85dbe138,07d13a8f,a8e962af,c449f783,27c07bd6,1f868fdd,21ddcdc9,a458ea53,7eee76d1,,32c7478e,9af06ad9,9d93af03,cdfe5ab7,"},
+    {"key": "example28", "data": "1,0,20,2,2,7188,170,2,3,24,0,2,0,2,68fd1e64,38a947a1,ee6e4611,30d9fc77,4cf72387,7e0ccccf,bf9d4f90,0b153874,a73ee510,b7c4dad5,81cae03e,5332e3fb,d413ef3e,07d13a8f,a6d97bf2,ec676ace,3486227d,02e8d897,,,b055c31b,,3a171ecb,ae2cd100,,,"},
+    {"key": "example29", "data": "1,0,78,2,15,4311,85,4,18,230,0,3,,15,68fd1e64,1287a654,5ed035c9,5b5365b2,4cf72387,6f6d9be8,1b1aa9ea,0b153874,a73ee510,c3e69838,7a3651f5,df8b1dea,95bc260c,b28479f6,ced5be3a,4cc0abe4,e5ba7672,df00d249,,,f520f961,,32c7478e,27b60b01,,,"},
+    {"key": "example30", "data": "1,3,0,4,13,224,28,3,35,27,1,1,,13,05db9164,90081f33,993f507e,14a74146,25c83c98,13718bbd,dc7659bd,0b153874,a73ee510,03e48276,e51ddf94,18fe7085,3516f6e6,64c94865,98995c3b,8c48eb08,e5ba7672,7181ccc8,,,2ed6b316,,3a171ecb,abf08f1b,,,"},
+    {"key": "example31", "data": "1,,277,,3,7318,24,6,3,98,,1,,3,8cf07265,9adf4cf9,2e76fb61,0b1ad9da,4cf72387,fe6b92e5,75dcaaca,0b153874,a73ee510,3b08e48b,8aabdae8,9886a0a7,edcf17ce,07d13a8f,2aaebd23,338c0d09,e5ba7672,c7dbecd5,,,60d2d691,,3a171ecb,90b6276f,,,"},
+    {"key": "example32", "data": "0,,-1,,,4956,,0,37,97,,0,,,be589b51,4c2bc594,d032c263,c18be181,25c83c98,fe6b92e5,aa0d873c,0b153874,a73ee510,3b08e48b,868744ab,dfbb09fb,9dfda2b9,8ceecbc8,7ac43a46,84898b2a,776ce399,bc48b783,,,0014c32a,,55dd3565,3b183c5c,,,"},
+    {"key": "example33", "data": "0,1,0,1,,1427,3,16,11,50,0,2,1,,05db9164,26a88120,615e3e4e,2788fed8,4cf72387,7e0ccccf,3f4ec687,0b153874,a73ee510,0e9ead52,c4adf918,f5d19c1c,85dbe138,07d13a8f,24ff9452,1034ac0d,3486227d,b486119d,,,63580fba,,32c7478e,2a90c749,,,"},
+    {"key": "example34", "data": "0,4,0,55,8,859,13,4,12,13,1,1,,8,05db9164,e5fb1af3,4b644986,5dbf0cc5,25c83c98,,cc5ed2f1,0b153874,a73ee510,3b08e48b,facf05cc,6d89b6a5,9f16a973,cfef1c29,1e744fde,2bd32e5c,776ce399,13145934,21ddcdc9,a458ea53,1419c3fc,,32c7478e,8ecc176a,e8b83407,a70a038a,"},
+    {"key": "example35", "data": "1,1,259,1,1,5,1,6,1,1,1,3,,1,05db9164,f3b07830,ad981000,f96c819d,25c83c98,,df5c2d18,0b153874,a73ee510,8aef4905,a7b606c4,b912be9f,eae197fd,b28479f6,d27eed0e,b8b09fe6,e5ba7672,048d01f4,,,08ae854d,,32c7478e,c657e6e5,,,"},
+    {"key": "example36", "data": "1,0,127,1,3,1683,19,26,17,475,0,9,0,3,05db9164,8947f767,11c9d79e,52a787c8,4cf72387,fbad5c96,18671b18,0b153874,a73ee510,ceb10289,77212bd7,79507c6b,7203f04e,07d13a8f,2c14c412,49013ffe,8efede7f,bd17c3da,f6a3e43b,a458ea53,35cd95c9,ad3062eb,c7dc6720,3fdb382b,010f6491,49d68486,"},
+    {"key": "example37", "data": "0,,1,,,23255,,0,1,73,,0,,,7e5c2ff4,d833535f,b00d1501,d16679b9,25c83c98,7e0ccccf,65c53f25,1f89b562,a73ee510,3b08e48b,ad2bc6f4,e0d76380,39ccb769,b28479f6,a733d362,1203a270,776ce399,281769c2,,,73d06dde,,32c7478e,aee52b6f,,,"},
+    {"key": "example38", "data": "0,6,-1,,,915,40,26,33,72,1,3,,,9a89b36c,4f25e98b,9042c4ea,343f8ed3,25c83c98,fbad5c96,27cc0b50,0b153874,a73ee510,f364a867,7671c62f,00750e7a,1fa0660e,b28479f6,df2f73e9,4f71659c,e5ba7672,bc5a0ff7,21ddcdc9,a458ea53,706ee322,c9d4222a,bcdee96c,990a118a,001f3601,47b6f269,"},
+    {"key": "example39", "data": "0,0,0,3,12,7308,97,2,21,90,0,1,,12,68fd1e64,9adf4cf9,723c059c,4c942c6d,4cf72387,7e0ccccf,ce4f7f55,0b153874,a73ee510,d7026747,38f692a7,ab60a748,6e5da64f,1adce6ef,808ff1bc,c23fc7ec,e5ba7672,2a93f7c8,,,5dc9a057,,32c7478e,90b6276f,,,"},
+    {"key": "example40", "data": "0,8,0,15,20,115,24,8,23,24,2,2,,20,5a9ed9b0,c66fca21,78171040,373c404a,25c83c98,,8ff6f5af,0b153874,a73ee510,5ba575e7,b5a9f90e,6766a7f0,949ea585,1adce6ef,8736735c,59974c9c,8efede7f,1304f63b,21ddcdc9,b1252a9d,07b2853e,,32c7478e,94bde4f2,010f6491,09b76f8d,"},
+    {"key": "example41", "data": "0,,38,2,4,3119,149,64,48,139,,6,6,4,05db9164,26a88120,d032c263,c18be181,4cf72387,fbad5c96,3f4ec687,1f89b562,a73ee510,726f00fd,c4adf918,dfbb09fb,85dbe138,07d13a8f,040ec437,84898b2a,8efede7f,57598e25,,,0014c32a,,32c7478e,3b183c5c,,,"},
+    {"key": "example42", "data": "1,88,319,,4,5,4,89,40,88,3,4,12,4,05db9164,08d6d899,333440d5,fc86bde0,25c83c98,fbad5c96,f00bddf8,0b153874,a73ee510,83ff688a,55795b33,1b0c8aa3,39795005,b28479f6,bffbd637,4a838997,8efede7f,bbf70d82,,,16e2e3b3,,32c7478e,d859b4dd,,,"},
+    {"key": "example43", "data": "0,,1,18,5,1683,80,38,5,95,,5,0,5,05db9164,09e68b86,aa8c1539,85dd697c,25c83c98,,2903ead3,0b153874,a73ee510,bcc8b4c6,a0a5e9d7,d8c29807,ee79db7b,1adce6ef,dcd06253,c64d548f,3486227d,63cdbb21,cf99e5de,a458ea53,5f957280,,32c7478e,1793a828,e8b83407,b7d9c3bc,"},
+    {"key": "example44", "data": "0,,27,,,112878,2106,0,2,95,,0,,,5a9ed9b0,38a947a1,2d8004c4,40ed41e5,25c83c98,7e0ccccf,4d9d55ae,5b392875,7cc72ec2,3b08e48b,55065437,ad972965,80dcea18,07d13a8f,c68ba31d,1206a8a1,d4bb7bd8,e96a7df2,,,54d8bb06,,3a171ecb,a415643d,,,"},
+    {"key": "example45", "data": "0,0,-1,,,4894,20,1,7,20,0,1,,,05db9164,4c2bc594,d032c263,c18be181,43b19349,7e0ccccf,7f52e00f,0b153874,a73ee510,48a4f593,bca79aeb,dfbb09fb,5218d824,8ceecbc8,7ac43a46,84898b2a,07c540c4,bc48b783,,,0014c32a,c9d4222a,3a171ecb,3b183c5c,,,"},
+    {"key": "example46", "data": "0,0,32,,1,9375,,0,37,18,0,0,0,1,05db9164,d833535f,ad4b77ff,d16679b9,25c83c98,7e0ccccf,9d547ce0,5b392875,a73ee510,3b08e48b,868a9e47,a2f4e8b5,fc5dea81,b28479f6,a733d362,89052618,3486227d,281769c2,,,d4703ebd,,32c7478e,aee52b6f,,,"},
+    {"key": "example47", "data": "0,,6,6,15,20213,507,7,42,360,,2,0,40,05db9164,0ca4b7d7,d032c263,c18be181,4cf72387,7e0ccccf,e9396c09,c8ddd494,a73ee510,3b08e48b,a0060bca,dfbb09fb,22d23aac,1adce6ef,9014f0f9,84898b2a,e5ba7672,c786d1ea,,,0014c32a,,32c7478e,3b183c5c,,,"},
+    {"key": "example48", "data": "0,31,17,2,11,290,23,31,23,65,2,2,,11,05db9164,4f25e98b,03280284,5214fda3,25c83c98,fbad5c96,0c41b6a1,0b153874,a73ee510,fa642b71,4ba74619,60bab41d,879fa878,07d13a8f,5be89da3,b6acbd10,e5ba7672,bc5a0ff7,fae651c5,a458ea53,3792328c,c0061c6d,423fab69,7a8e7ed6,001f3601,f159b6cb,"},
+    {"key": "example49", "data": "0,1,2382,13,4,40,4,69,3,609,1,11,0,4,05db9164,38a947a1,933cc823,b1c1e580,25c83c98,fe6b92e5,002fdf0c,1f89b562,a73ee510,61f70369,a4ea009a,2562cf3c,1e9339bc,b28479f6,f5bfabbd,03dee53f,e5ba7672,b3e92443,,,be661a75,,c7dc6720,67d37917,,,"},
+    {"key": "example50", "data": "0,0,190,,,1624,6,29,6,74,0,9,,,68fd1e64,c41a84c8,0a266224,759c4a2e,25c83c98,,804d2f11,0b153874,a73ee510,2860ede1,1aa6cf31,99dfd83a,3b03d76e,b28479f6,e3eb97c7,62b5674b,e5ba7672,5911fc7e,,,28ee216d,,32c7478e,590b856f,,,"},
+    {"key": "example51", "data": "0,0,19,8,9,1506,31,3,7,34,0,1,,9,05db9164,8947f767,100a3803,ad1b5124,30903e74,7e0ccccf,bb3b7ab9,c8ddd494,a73ee510,3b08e48b,90b202b5,d377c333,3a9dafb8,b28479f6,a473257f,68d2c2b9,8efede7f,bd17c3da,e51f040f,a458ea53,79c3f011,,bcdee96c,fe35ffe2,010f6491,987ea0be,"},
+    {"key": "example52", "data": "1,1,30,,3,116,31,1,3,3,1,1,,3,5a9ed9b0,38a947a1,0f90b6d6,16b922ed,25c83c98,fbad5c96,17b47bf9,0b153874,a73ee510,3b08e48b,76120d9d,ee2d3fdb,d4384424,07d13a8f,3046a70a,81371cbc,d4bb7bd8,1bae7658,,,37bad455,ad3062eb,3a171ecb,9d96bacb,,,"},
+    {"key": "example53", "data": "0,,6,2,3,2779,,0,3,13,,0,,3,fb174e6b,47e8ab98,b009d929,c7043c4b,384874ce,,646e7593,0b153874,a73ee510,3b08e48b,d05acfa9,3563ab62,969e14fd,1adce6ef,bfa6d08a,b688c8cc,8efede7f,eb4d3f8a,21ddcdc9,5840adea,2754aaf1,,55dd3565,3b183c5c,f55c04b6,491eeeef,"},
+    {"key": "example54", "data": "0,0,2,22,3,4687,242,6,6,183,0,1,4,3,05db9164,287130e0,c09cf4ef,bc8d1aa6,25c83c98,13718bbd,1919941b,37e4aa92,a73ee510,6c47047a,86c05043,c4bba41d,2ecea536,b28479f6,9efd8b77,ac2e5095,8efede7f,891589e7,2efde463,b1252a9d,dc4e98e3,,3a171ecb,ee42de86,e8b83407,a00829e6,"},
+    {"key": "example55", "data": "0,,55,16,7,1696,72,2,7,95,,2,,7,5bfa8ab5,89ddfee8,00e2b23c,10d65c35,25c83c98,7e0ccccf,ad3508b1,5b392875,a73ee510,fc3680e8,ad757a5a,f400e021,93b18cb5,1adce6ef,34cce7d2,9e87470c,e5ba7672,5bb2ec8e,7a45f7f2,a458ea53,a13d5eab,,423fab69,faf5d8b3,f0f449dd,a8cf207e,"},
+    {"key": "example56", "data": "0,6,0,28,0,31,0,6,0,0,1,1,,0,8cf07265,287130e0,c1ba4c5a,16fe249c,25c83c98,7e0ccccf,c1225605,985e3fcb,a73ee510,ede207dc,f29b9ed2,469027a9,7eaf6f1a,07d13a8f,10040656,8f13519e,e5ba7672,891589e7,6f3756eb,5840adea,f4095a39,,c7dc6720,1793a828,e8b83407,a475662f,"},
+    {"key": "example57", "data": "0,0,19,9,3,14414,1353,3,1,362,0,1,0,3,be589b51,09e68b86,4bee8a47,5031d726,25c83c98,7e0ccccf,197b4575,322e63df,a73ee510,6c47047a,48876b80,6f95f18b,e40e52ae,07d13a8f,36721ddc,3d66d729,8efede7f,5aed7436,21ddcdc9,a458ea53,3c1a8dd8,,3a171ecb,3fdb382b,b9266ff0,49d68486,"},
+    {"key": "example58", "data": "0,,18,12,8,8965,44,2,12,57,,2,0,8,05db9164,d7988e72,5eee7056,c1a3acf5,afcf7897,,96825c8f,0b153874,a73ee510,bc283a64,2df02cf1,f71904ea,03232503,b28479f6,c8389df7,0db58836,07c540c4,0f2f9850,5fd56cf9,b1252a9d,96725293,,32c7478e,2702453c,8b8de563,303cea07,"},
+    {"key": "example59", "data": "0,,1,2,0,177674,,0,3,2,,0,0,1,87552397,207b2d81,6e136288,4f938621,25c83c98,7e0ccccf,8025502e,6c41e35e,7cc72ec2,4072f40f,29e4ad33,64ddde07,80467802,07d13a8f,0bf0feff,0c41b634,e5ba7672,fa0643ee,21ddcdc9,b1252a9d,b4031b95,,3a171ecb,a81956df,001f3601,b1262ddd,"},
+    {"key": "example60", "data": "0,,0,20,0,4412,855,0,4,522,,0,,5,05db9164,58e67aaf,54e3c628,9725d851,25c83c98,7e0ccccf,5b18f3d9,0b153874,a73ee510,ad0b97fb,720446f5,92409ea2,034e5f3b,07d13a8f,10935a85,05c5cfbe,d4bb7bd8,c21c3e4c,6f62a118,b1252a9d,54d0b766,,c7dc6720,2913df0f,9b3e8820,bc7f21c2,"},
+    {"key": "example61", "data": "0,,56,2,,,,0,0,1,,0,,,5a9ed9b0,8084ee93,02cf9876,c18be181,25c83c98,fbad5c96,af0809a5,5b392875,7cc72ec2,3b08e48b,9e12e146,8fe001f4,025225f2,b28479f6,16d2748c,36103458,2005abd1,003d4f4f,,,e587c466,,be7c41b4,3b183c5c,,,"},
+    {"key": "example62", "data": "0,,76,5,,46200,,,7,,,,0,,68fd1e64,287130e0,7555338e,e161fae2,25c83c98,7e0ccccf,ce17d537,0b153874,7cc72ec2,ed111662,5b225578,f34e8f6a,d1be539d,07d13a8f,10040656,8ec308fc,3486227d,891589e7,21ddcdc9,5840adea,182fdd1a,,c7dc6720,6c1cdd05,ea9a246c,1219b447,"},
+		{"key": "example63", "data": "0,39,8,42,32,27,33,39,24,32,1,1,,32,05db9164,73a46ff0,844ce0a4,9bb11257,4cf72387,7e0ccccf,ff3f3dda,0fb392dd,a73ee510,3b08e48b,e2217f93,1d0b8187,da9ee8bd,1adce6ef,d57668e2,4372eb4b,e5ba7672,da507f45,21ddcdc9,5840adea,ecb5cd6f,,32c7478e,6dbd889f,ea9a246c,33ced911,"},
+		{"key": "example64", "data": "1,,0,55,5,14477,,0,5,1,,0,,5,05db9164,09e68b86,2beedeb2,c59396e7,25c83c98,7e0ccccf,a972360e,0b153874,7cc72ec2,acf0058d,9e511730,8f70e33a,04e4a7e0,64c94865,91126f30,6df0eed9,e5ba7672,5aed7436,55dd3565,5840adea,c412f773,,3a171ecb,3fdb382b,e8b83407,ccc71a58,"},
+		{"key": "example65", "data": "0,1,1,5,8,7,8,1,8,8,1,1,0,8,05db9164,8e4f887c,,,25c83c98,13718bbd,47802627,5b9f3341,a73ee510,fbbf2c95,42e01668,,79db54f6,07d13a8f,b708086d,,d4bb7bd8,4b340164,,,,,3a171ecb,,,,"},
+		{"key": "example66", "data": "1,,1,7,3,10087,67,6,3,57,,3,,3,05db9164,d7988e72,fb535e16,6fe3d332,25c83c98,13718bbd,2829f187,66f29b89,a73ee510,e034d733,3a9c7259,8e86918c,0d8d4492,07d13a8f,194c42a4,57236df8,e5ba7672,0f2f9850,c27239bd,5840adea,841f2712,,3a171ecb,1793a828,e8b83407,b820b6c5,"},
+		{"key": "example67", "data": "0,7,1,40,,1418,23,147,0,7,0,4,0,,68fd1e64,80e26c9b,ba1947d0,85dd697c,25c83c98,7e0ccccf,16401b7d,a61cc0ef,a73ee510,3b08e48b,20ec800a,34a238e0,18a5e4b8,b28479f6,a785131a,da441c7e,e5ba7672,005c6740,21ddcdc9,5840adea,8717ea07,,32c7478e,1793a828,e8b83407,b9809574,"},
+		{"key": "example68", "data": "0,,0,10,2,3545,,0,2,3,,0,,2,be589b51,e5fb1af3,50808b4e,39ec3719,25c83c98,13718bbd,316949b7,5b392875,a73ee510,3b08e48b,d51f40d7,75ef3efe,4eb5dabc,07d13a8f,b5de5956,d5cb04e4,776ce399,13145934,21ddcdc9,b1252a9d,1d4696ef,,32c7478e,39fe175c,e8b83407,1c7f8927,"},
+		{"key": "example69", "data": "0,5,51,,5,457,5,5,7,11,1,1,1,5,ae82ea21,9e5ce894,e1120103,13508380,25c83c98,7e0ccccf,6855ef53,0b153874,a73ee510,175d6c71,b7094596,d19a1cc6,1f9d2c38,07d13a8f,8cf98699,e58b9a62,3486227d,a5bb7b8a,1d1eb838,a458ea53,2eb5be02,ad3062eb,c7dc6720,45ab94c8,ea9a246c,c84c4aec,"},
+		{"key": "example70", "data": "0,0,0,15,2,20112,305,1,43,228,0,1,,30,05db9164,0a519c5c,02cf9876,c18be181,25c83c98,7e0ccccf,fe4e75fa,0b153874,a73ee510,6aea41c7,8f4f8f83,8fe001f4,8828a59c,07d13a8f,4ac81a35,36103458,d4bb7bd8,416e8695,,,e587c466,,93bad2c0,3b183c5c,,,"},
+		{"key": "example71", "data": "0,0,1,4,3,1689,184,12,46,53,0,1,,3,8cf07265,207b2d81,d6be853a,4842a03d,384874ce,fe6b92e5,209d1929,5b392875,a73ee510,9eff685e,87fe3e10,dadde5ca,3bd6c21d,b28479f6,3c767806,6077db2c,e5ba7672,395856b0,21ddcdc9,a458ea53,ae7b2d98,,32c7478e,b8942a02,001f3601,4a6648b5,"},
+		{"key": "example72", "data": "0,4,0,89,4,486,5,4,7,4,1,1,,4,05db9164,8947f767,22e8ec23,f8faa363,43b19349,7e0ccccf,45607029,51d76abe,a73ee510,0ada1061,2b9f131d,18e1f914,aca10c14,1adce6ef,ba8b8b16,226c87e7,d4bb7bd8,bd17c3da,656485cf,a458ea53,0443b252,,bcdee96c,c73755d6,e8b83407,c23979db,"},
+		{"key": "example73", "data": "1,,39,2,1,3343,1,17,2,1,,1,,1,05db9164,5a88f1d5,16424a73,08694bce,30903e74,fe6b92e5,12c61956,5b392875,a73ee510,52486df2,94d2aad8,c76cdf17,f23a3825,b28479f6,e842876b,ffe60785,e5ba7672,1adff463,,,5f0fcebd,ad3062eb,3a171ecb,392bf8f1,,,"},
+		{"key": "example74", "data": "0,,2,1,2,8036,164,18,13,367,,2,0,2,68fd1e64,287130e0,54597e12,12c911a7,384874ce,13718bbd,48b70cb6,985e3fcb,a73ee510,6123dced,1736789a,8c51bef7,1a347339,b28479f6,9efd8b77,f8f7edbf,3486227d,891589e7,99f90f6d,5840adea,72b1423f,,32c7478e,af062947,ea9a246c,b4a4615f,"},
+		{"key": "example75", "data": "0,,-1,,,71140,142,0,7,63,,0,,,8cf07265,4c2bc594,d032c263,c18be181,25c83c98,fbad5c96,5d859d57,0b153874,7cc72ec2,cd481139,00adbfbb,dfbb09fb,d4b85d8d,8ceecbc8,7ac43a46,84898b2a,07c540c4,bc48b783,,,0014c32a,,3a171ecb,3b183c5c,,,"},
+		{"key": "example76", "data": "1,,0,11,3,16184,125,2,3,103,,2,0,3,05db9164,f0cf0024,6f67f7e5,41274cd7,25c83c98,,94a113a4,0b153874,a73ee510,4ddb41b1,f47e21eb,623049e6,4f3f2bb1,b28479f6,e6c5b5cd,c92f3b61,07c540c4,b04e4670,21ddcdc9,5840adea,60f6221e,,32c7478e,43f13e8b,ea9a246c,731c3655,"},
+		{"key": "example77", "data": "0,,140,2,2,,,0,2,2,,0,,2,5bfa8ab5,38a947a1,,,25c83c98,7e0ccccf,88002ee1,64523cfa,7cc72ec2,3b08e48b,f1b78ab4,,6e5da64f,07d13a8f,c2b7aaa6,,2005abd1,659bdb63,,,,ad3062eb,32c7478e,,,,"},
+		{"key": "example78", "data": "0,2,0,,7,443,37,7,34,282,1,4,7,7,3c9d8785,38a947a1,4470baf4,8c8a4c47,43b19349,fbad5c96,282b88fc,0b153874,a73ee510,0f1a2599,ea26a3ee,bb669e25,0e5bc979,b28479f6,547b8c62,2b2ce127,8efede7f,b133fcd4,,,2b796e4a,,32c7478e,8d365d3b,,,"},
+		{"key": "example79", "data": "1,2,1,60,75,61,121,52,39,248,1,8,1,77,05db9164,942f9a8d,ab4a038c,fea9881c,4cf72387,7e0ccccf,3f4ec687,0b153874,a73ee510,726f00fd,c4adf918,2b7b1137,85dbe138,1adce6ef,ae97ecc3,2dbf1d23,8efede7f,1f868fdd,f44bef3c,a458ea53,8adfc28d,,bcdee96c,3fdb382b,9d93af03,49d68486,"},
+		{"key": "example80", "data": "0,1,34,2,42,328,44,15,49,58,1,9,0,42,05db9164,2c16a946,a65db9fb,9f43a1b5,25c83c98,,1d794a16,5b392875,a73ee510,ed086ca2,4c9e8313,28156fd4,67b031b4,b28479f6,3628a186,87140baa,e5ba7672,e4ca448c,,,67bb5322,,32c7478e,9117a34a,,,"},
+		{"key": "example81", "data": "0,,0,,,14919,,0,0,0,,0,,,05db9164,4c2bc594,d032c263,c18be181,384874ce,fe6b92e5,b7c924a4,64523cfa,a73ee510,3b08e48b,2cc0193e,dfbb09fb,433f9499,8ceecbc8,7ac43a46,84898b2a,1e88c74f,bc48b783,,,0014c32a,,3a171ecb,3b183c5c,,,"},
+		{"key": "example82", "data": "0,,1,,1,20216,66,1,1,1,,1,,1,68fd1e64,ae46a29d,c68767c4,f922efad,4cf72387,fbad5c96,5fd3419b,0b153874,a73ee510,972359d0,efc5e2cf,9e62f74b,c7176043,8ceecbc8,14d0a096,e2e2fcd9,d4bb7bd8,a573334c,,,3c35236f,,55dd3565,b34f3128,,,"},
+		{"key": "example83", "data": "0,13,13,80,32,378,115,15,37,57,1,2,,48,5a9ed9b0,4f25e98b,b393df87,c3fecae9,25c83c98,7e0ccccf,5192dba2,0b153874,a73ee510,f1317066,aaa08406,8d33fe00,6665daff,8ceecbc8,5525889d,0fd4fbad,e5ba7672,9e4517be,3014a4b1,5840adea,572bdde8,ad3062eb,32c7478e,9fa3e01a,001f3601,d9bcfc08,"},
+		{"key": "example84", "data": "0,,0,91,4,293,220,35,4,66,,5,1,4,05db9164,58e67aaf,1100da42,5024ec42,4cf72387,7e0ccccf,964d1fdd,37e4aa92,a73ee510,f007059b,59cd5ae7,a9d84aa8,8b216f7b,1adce6ef,d002b6d9,3b35c4c2,27c07bd6,c21c3e4c,21ddcdc9,a458ea53,e9ab8737,,423fab69,2c720b71,9b3e8820,47d6d9aa,"},
+		{"key": "example85", "data": "0,,-1,,,5654,,0,4,15,,0,0,,05db9164,09e68b86,d8d5b6da,bd97eb27,43b19349,fbad5c96,4b219154,1f89b562,a73ee510,7259dc52,f25fe7e9,a6cbd3ab,dd183b4c,07d13a8f,36721ddc,cd220c47,3486227d,5aed7436,338f20de,a458ea53,7e08e349,,32c7478e,661d665d,9d93af03,ee88160c,"},
+		{"key": "example86", "data": "0,0,53,,10,6550,98,34,11,349,0,9,,10,05db9164,207b2d81,8bd78c57,394ee067,25c83c98,6f6d9be8,283d5555,0b153874,a73ee510,3b08e48b,3d5fb018,e5f6b330,94172618,07d13a8f,0bf0feff,402a9036,e5ba7672,fa0643ee,21ddcdc9,b1252a9d,0094bc78,,32c7478e,29ece3ed,001f3601,402185f3,"},
+		{"key": "example87", "data": "1,6,-1,10,6,0,0,6,9,9,1,1,,0,05db9164,09e68b86,be7504db,2f91f54d,25c83c98,7e0ccccf,26a81064,5b392875,a73ee510,dcbc7c2b,9e511730,1205ef20,04e4a7e0,07d13a8f,36721ddc,9f064e1a,e5ba7672,5aed7436,1d04f4a4,a458ea53,c3444bea,,3a171ecb,3fdb382b,e8b83407,49d68486,"},
+		{"key": "example88", "data": "1,,0,1,3,2901,,0,3,9,,0,,3,05db9164,762b9a6f,b559a4bf,c7663617,25c83c98,,1f3230eb,1f89b562,a73ee510,687c8b31,05766aa3,051dfd94,b94e5df6,07d13a8f,aa600b94,dad720cc,1e88c74f,01890ebf,,,05afbcbe,,32c7478e,5a1a48d4,,,"},
+		{"key": "example89", "data": "0,1,23,1,14,56,15,10,35,161,1,2,,15,be589b51,38a947a1,,,25c83c98,,ba0ca6c5,0b153874,a73ee510,56ae5fb0,7ca01a9d,,97d749c9,07d13a8f,bbe9d023,,e5ba7672,f9c50fb1,,,,,32c7478e,,,,"},
+		{"key": "example90", "data": "0,,1,75,22,5912,181,2,22,68,,1,,22,5bfa8ab5,09e68b86,aa8c1539,85dd697c,4cf72387,,845fb196,5b392875,a73ee510,6997b535,911b463c,d8c29807,0c9cc756,8ceecbc8,d2f03b75,c64d548f,07c540c4,63cdbb21,cf99e5de,5840adea,5f957280,,32c7478e,1793a828,e8b83407,b7d9c3bc,"},
+		{"key": "example91", "data": "1,,127,2,23,15837,23,4,24,23,,1,,23,68fd1e64,d4be07ad,d8fd96ac,9581b80f,25c83c98,fe6b92e5,ea3f0578,0b153874,a73ee510,c514dac9,42880796,2af3c6b9,6e1e209e,07d13a8f,1936a526,94df6876,e5ba7672,cbae5931,1f4a1e60,5840adea,115892da,,c7dc6720,b2f178a3,001f3601,938732a0,"},
+		{"key": "example92", "data": "0,,-1,,,23204,,0,32,38,,0,,,68fd1e64,4c2bc594,d032c263,c18be181,25c83c98,fe6b92e5,3d63f4e6,0b153874,a73ee510,376bbe93,af6a4ffc,dfbb09fb,2a1579a2,1adce6ef,ae0c3875,84898b2a,8efede7f,15a36060,,,0014c32a,,423fab69,3b183c5c,,,"},
+		{"key": "example93", "data": "0,,-1,,,681386,,,11,,,,,,05db9164,4c2bc594,d032c263,c18be181,43b19349,7e0ccccf,1554a783,0b153874,7cc72ec2,7636f6c8,b7bb7a17,dfbb09fb,73e186f6,8ceecbc8,7ac43a46,84898b2a,e5ba7672,bc48b783,,,0014c32a,c9d4222a,3a171ecb,3b183c5c,,,"},
+		{"key": "example94", "data": "0,,1,64,11,1494,86,15,13,231,,4,,11,05db9164,8ab240be,f1fce64b,fbe29d7d,25c83c98,7e0ccccf,22eefdf2,0b153874,a73ee510,994730d3,974aaa98,51874c73,7889204d,07d13a8f,e7dd0bfc,d2f669a0,e5ba7672,ca533012,21ddcdc9,5840adea,30c1399f,,32c7478e,6856d4e9,b9266ff0,5f3ed85b,"},
+		{"key": "example95", "data": "0,,1,1,,2903,2,3,0,2,,1,,,05db9164,80e26c9b,aa8c1539,85dd697c,43b19349,fbad5c96,9c7e110d,0b153874,a73ee510,3b08e48b,3204e37d,d8c29807,8b4df3b4,1adce6ef,0f942372,c64d548f,07c540c4,005c6740,21ddcdc9,5840adea,5f957280,,3a171ecb,1793a828,e8b83407,9904c656,"},
+		{"key": "example96", "data": "0,,105,1,1,171519,,0,1,13,,0,,1,05db9164,38a947a1,145f2f75,82a61820,4cf72387,7e0ccccf,9d99aaa3,5b392875,7cc72ec2,2ce2764d,d93e6010,7161e106,4e8bba73,1adce6ef,402ce7c7,bb6d240e,07c540c4,e96a7df2,,,5fe17899,,423fab69,cafb4e4d,,,"},
+		{"key": "example97", "data": "0,,35,90,2,41230,352,0,5,44,,0,,5,05db9164,c5fe64d9,4e14c1ea,cb4db510,25c83c98,7e0ccccf,77e91f62,1f89b562,a73ee510,1a88cf9b,7defe259,4ec1453d,11fa2c12,b28479f6,543c0413,9cdd6d81,e5ba7672,c235abed,4632bcdc,b1252a9d,4bab48f8,,c7dc6720,3fdb382b,ea9a246c,49d68486,"},
+		{"key": "example98", "data": "0,,0,78,1,15835,220,1,1,1,,1,,1,05db9164,09e68b86,adeee51f,c9add53e,25c83c98,fe6b92e5,a255dd63,0b153874,a73ee510,4effc25c,d13e1160,75762013,45820f61,07d13a8f,36721ddc,afb79297,d4bb7bd8,5aed7436,363a4009,b1252a9d,aaf7d15f,,3a171ecb,1793a828,e8b83407,d94377ca,"},
+		{"key": "example99", "data": "0,,1715,,1,7545,19,15,1,87,,7,0,1,5a9ed9b0,999aae35,79bc99b4,e5e453f3,25c83c98,7e0ccccf,e13ff1c4,1f89b562,a73ee510,fbbf2c95,15cd287d,424e28fe,86c79eb0,243a4e68,39a6addf,a6a69939,3486227d,63aa00dd,,,424af181,,3a171ecb,869caea3,,,"},
+		{"key": "example100", "data": "0,0,48,3,2,14482,377,8,2,283,0,2,,2,05db9164,09e68b86,b3bebfb8,bc93ba48,43b19349,,46319bad,0b153874,a73ee510,14e49183,c92f4124,49fe61e3,c576c612,b28479f6,52baadf5,dc12c38d,e5ba7672,5aed7436,21ddcdc9,b1252a9d,0421c25a,,32c7478e,a212d987,e8b83407,f82917a3"}
   ]
 }
diff --git a/tests/endtoendtests/OdpSampleNetwork/src/com/example/odpsamplenetwork/SampleHandler.java b/tests/endtoendtests/OdpSampleNetwork/src/com/example/odpsamplenetwork/SampleHandler.java
index d0e255e5..1bc4306e 100644
--- a/tests/endtoendtests/OdpSampleNetwork/src/com/example/odpsamplenetwork/SampleHandler.java
+++ b/tests/endtoendtests/OdpSampleNetwork/src/com/example/odpsamplenetwork/SampleHandler.java
@@ -34,8 +34,8 @@ import android.adservices.ondevicepersonalization.RenderInput;
 import android.adservices.ondevicepersonalization.RenderOutput;
 import android.adservices.ondevicepersonalization.RenderingConfig;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
-import android.adservices.ondevicepersonalization.TrainingExampleInput;
-import android.adservices.ondevicepersonalization.TrainingExampleOutput;
+import android.adservices.ondevicepersonalization.TrainingExamplesInput;
+import android.adservices.ondevicepersonalization.TrainingExamplesOutput;
 import android.adservices.ondevicepersonalization.TrainingInterval;
 import android.adservices.ondevicepersonalization.UserData;
 import android.content.ContentValues;
@@ -70,10 +70,12 @@ import java.io.IOException;
 import java.io.StringReader;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Random;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadFactory;
 import java.util.function.Consumer;
@@ -145,42 +147,37 @@ public class SampleHandler implements IsolatedWorker {
         consumer.accept(downloadResult);
     }
 
-    @Override public void onExecute(
-            @NonNull ExecuteInput input,
-            @NonNull Consumer<ExecuteOutput> consumer
-    ) {
+    @Override
+    public void onExecute(@NonNull ExecuteInput input, @NonNull Consumer<ExecuteOutput> consumer) {
         Log.d(TAG, "onExecute() started.");
         sBackgroundExecutor.execute(() -> handleOnExecute(input, consumer));
     }
 
-    @Override public void onTrainingExample(
-            @NonNull TrainingExampleInput input,
-            @NonNull Consumer<TrainingExampleOutput> consumer) {
-        Log.d(TAG, "onTrainingExample() started.");
-        sBackgroundExecutor.execute(() -> handleOnTrainingExample(input, consumer));
+    @Override
+    public void onTrainingExamples(
+            @NonNull TrainingExamplesInput input,
+            @NonNull Consumer<TrainingExamplesOutput> consumer) {
+        Log.d(TAG, "onTrainingExamples() started.");
+        sBackgroundExecutor.execute(() -> handleOnTrainingExamples(input, consumer));
     }
 
-    @Override public void onRender(
-            @NonNull RenderInput input,
-            @NonNull Consumer<RenderOutput> consumer
-    ) {
+    @Override
+    public void onRender(@NonNull RenderInput input, @NonNull Consumer<RenderOutput> consumer) {
         Log.d(TAG, "onRender() started.");
         sBackgroundExecutor.execute(() -> handleOnRender(input, consumer));
     }
 
-    @Override public void onEvent(
-            @NonNull EventInput input,
-            @NonNull Consumer<EventOutput> consumer) {
+    @Override
+    public void onEvent(@NonNull EventInput input, @NonNull Consumer<EventOutput> consumer) {
         Log.d(TAG, "onEvent() started.");
-        sBackgroundExecutor.execute(
-                () -> handleOnWebViewEvent(input, consumer));
+        sBackgroundExecutor.execute(() -> handleOnWebViewEvent(input, consumer));
     }
 
     private ListenableFuture<List<Ad>> readAds(KeyValueStore remoteData) {
         Log.d(TAG, "readAds() called.");
         try {
             ArrayList<Ad> ads = new ArrayList<>();
-            for (var key: remoteData.keySet()) {
+            for (var key : remoteData.keySet()) {
                 if (!key.startsWith("ad")) {
                     continue;
                 }
@@ -240,7 +237,7 @@ public class SampleHandler implements IsolatedWorker {
         String requestKeyword = "";
         if (input != null && input.getAppParams() != null
                 && input.getAppParams().getString("keyword") != null) {
-            requestKeyword = input.getAppParams().getString("keyword");
+            requestKeyword = input.getAppParams().getString("keyword").toLowerCase().strip();
         }
 
         List<Ad> result = new ArrayList<>();
@@ -282,73 +279,55 @@ public class SampleHandler implements IsolatedWorker {
                 .build();
     }
 
-    static Feature convertStringToFeature(String value) {
+    private static Feature convertStringToFeature(String value) {
         BytesList.Builder bytesListBuilder = BytesList.newBuilder();
         String nonNullValue = Strings.nullToEmpty(value);
         bytesListBuilder.addValue(ByteString.copyFromUtf8(nonNullValue));
         return Feature.newBuilder().setBytesList(bytesListBuilder.build()).build();
     }
 
-    static Feature convertLongToFeature(long value) {
+    private static Feature convertLongToFeature(String value) {
+        long lValue = value.isEmpty() ? 0L : Long.parseLong(value);
         return Feature.newBuilder()
-            .setInt64List(Int64List.newBuilder().addValue(value).build())
+            .setInt64List(Int64List.newBuilder().addValue(lValue).build())
             .build();
     }
 
-    private void handleOnTrainingExample(
-            @NonNull TrainingExampleInput input,
-            @NonNull Consumer<TrainingExampleOutput> consumer) {
+    private static Example convertToExample(String serializedExample) {
+        String[] splitExample = serializedExample.split(",", -1);
         Features.Builder featuresBuilder = Features.newBuilder();
+        featuresBuilder.putFeature("clicked", convertLongToFeature(splitExample[0]));
+
+        int count = 1;
+        for (; count < 14; count++) {
+            featuresBuilder.putFeature(
+                    String.format("int-feature-%d", count),
+                    convertLongToFeature(splitExample[count]));
+        }
+        for (; count < 40; count++) {
+            featuresBuilder.putFeature(
+                    String.format("categorical-feature-%d", count),
+                    convertStringToFeature(splitExample[count]));
+        }
+        return Example.newBuilder().setFeatures(featuresBuilder.build()).build();
+    }
 
-        featuresBuilder.putFeature("int-feature-1", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-2", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-3", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-4", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-5", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-6", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-7", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-8", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-9", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-10", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-11", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-12", convertLongToFeature(0L));
-        featuresBuilder.putFeature("int-feature-13", convertLongToFeature(0L));
-
-        featuresBuilder.putFeature("categorical-feature-14", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-15", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-16", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-17", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-18", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-19", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-20", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-21", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-22", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-23", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-24", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-25", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-26", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-27", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-28", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-29", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-30", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-31", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-32", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-33", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-34", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-35", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-36", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-37", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-38", convertStringToFeature(""));
-        featuresBuilder.putFeature("categorical-feature-39", convertStringToFeature(""));
-
-        featuresBuilder.putFeature("clicked", convertLongToFeature(1L));
-
-        Example example = Example.newBuilder().setFeatures(featuresBuilder.build()).build();
-        TrainingExampleOutput result = new TrainingExampleOutput
-                .Builder()
-                .addTrainingExample(example.toByteArray())
-                .addResumptionToken("token1".getBytes()).build();
-        consumer.accept(result);
+    private void handleOnTrainingExamples(
+            @NonNull TrainingExamplesInput input,
+            @NonNull Consumer<TrainingExamplesOutput> consumer) {
+        TrainingExamplesOutput.Builder resultBuilder = new TrainingExamplesOutput.Builder();
+        Random rand = new Random();
+        int numExample = rand.nextInt(10) + 1;
+        Log.d(TAG, String.format("onTrainingExample() generates %d examples.", numExample));
+        for (int count = 0; count < numExample; count++) {
+            Example example = convertToExample(
+                    new String(mRemoteData.get(String.format("example%d", rand.nextInt(100) + 1)),
+                    StandardCharsets.UTF_8));
+            resultBuilder.addTrainingExample(
+                    example.toByteArray()).addResumptionToken(
+                            String.format("token%d", count).getBytes());
+        }
+        consumer.accept(resultBuilder.build());
     }
 
     private void handleOnExecute(
@@ -469,7 +448,8 @@ public class SampleHandler implements IsolatedWorker {
         }
         long now = System.currentTimeMillis();
         List<EventLogRecord> logRecords = mLogReader.getJoinedEvents(
-                now - 24 * 60 * 60 * 1000, now);
+                Instant.ofEpochMilli(now - 24 * 60 * 60 * 1000),
+                Instant.ofEpochMilli(now));
         EventLogRecord found = null;
         // Attribute conversion to most recent impression or click.
         for (EventLogRecord ev : logRecords) {
@@ -646,7 +626,7 @@ public class SampleHandler implements IsolatedWorker {
                     if (ad != null && !isBlockedAd(ad)) {
                         filteredKeys.add(key);
                     }
-                } else if (key.startsWith("template")) {
+                } else if (key.startsWith("template") || key.startsWith("example")) {
                     filteredKeys.add(key);
                 }
             }
diff --git a/tests/federatedcomputetests/AndroidManifest.xml b/tests/federatedcomputetests/AndroidManifest.xml
index 9767a489..09591632 100644
--- a/tests/federatedcomputetests/AndroidManifest.xml
+++ b/tests/federatedcomputetests/AndroidManifest.xml
@@ -32,6 +32,10 @@
                  android:exported="false"
                  android:permission="android.permission.BIND_JOB_SERVICE">
         </service>
+        <service android:name="com.android.federatedcompute.services.encryption.BackgroundKeyFetchJobService"
+            android:exported="false"
+            android:permission="android.permission.BIND_JOB_SERVICE">
+        </service>
         <service android:name="com.android.federatedcompute.services.training.IsolatedTrainingService"
                  android:isolatedProcess="true" android:exported="false" >
         </service>
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImplTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImplTest.java
index 43e98e72..ecdeddab 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImplTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/FederatedComputeManagingServiceImplTest.java
@@ -17,7 +17,9 @@
 package com.android.federatedcompute.services;
 
 import static org.junit.Assert.assertNotNull;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 
 import android.content.Intent;
 import android.os.IBinder;
@@ -25,6 +27,7 @@ import android.os.IBinder;
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.dx.mockito.inline.extended.ExtendedMockito;
+import com.android.federatedcompute.services.encryption.BackgroundKeyFetchJobService;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -39,7 +42,10 @@ public final class FederatedComputeManagingServiceImplTest {
 
     @Test
     public void testBindableFederatedComputeService() {
-        MockitoSession session = ExtendedMockito.mockitoSession().startMocking();
+        MockitoSession session = ExtendedMockito.mockitoSession()
+                .spyStatic(BackgroundKeyFetchJobService.class).startMocking();
+        ExtendedMockito.doReturn(true)
+                .when(() -> BackgroundKeyFetchJobService.scheduleJobIfNeeded(any(), any()));
         try {
             FederatedComputeManagingServiceImpl spyFcpService =
                     spy(new FederatedComputeManagingServiceImpl());
@@ -49,6 +55,8 @@ public final class FederatedComputeManagingServiceImplTest {
                             ApplicationProvider.getApplicationContext(),
                             FederatedComputeManagingServiceImpl.class);
             IBinder binder = spyFcpService.onBind(intent);
+            ExtendedMockito.verify(() -> BackgroundKeyFetchJobService.scheduleJobIfNeeded(
+                    any(), any()), times(1));
             assertNotNull(binder);
         } finally {
             session.finishMocking();
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 10c362c3..90340bf9 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/common/PhFlagsTest.java
@@ -17,20 +17,20 @@
 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.PhFlags.FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL;
 import static com.android.federatedcompute.services.common.PhFlags.KEY_FEDERATED_COMPUTE_KILL_SWITCH;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import android.provider.DeviceConfig;
 
-import androidx.test.runner.AndroidJUnit4;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
 
 /** Unit tests for {@link com.android.ondevicepersonalization.service.PhFlags} */
-@RunWith(AndroidJUnit4.class)
+@RunWith(JUnit4.class)
 public class PhFlagsTest {
     /** Get necessary permissions to access Setting.Config API and set up context */
     @Before
@@ -60,4 +60,18 @@ public class PhFlagsTest {
         Flags phFlags = FlagsFactory.getFlags();
         assertThat(phFlags.getGlobalKillSwitch()).isEqualTo(phOverridingValue);
     }
+
+    @Test
+    public void testGetEncryptionKeyFetchUrl() {
+        // Now overriding with the value from PH.
+        String overrideUrl = "https://real-coordinator/v1alpha/publicKeys";
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ON_DEVICE_PERSONALIZATION,
+                FEDERATED_COMPUTATION_ENCRYPTION_KEY_DOWNLOAD_URL,
+                overrideUrl,
+                /* makeDefault= */ false);
+
+        Flags phFlags = FlagsFactory.getFlags();
+        assertThat(phFlags.getEncryptionKeyFetchUrl()).isEqualTo(overrideUrl);
+    }
 }
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 6daff32d..b6f1bbc2 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/data/FederatedTrainingTaskDaoTest.java
@@ -59,9 +59,8 @@ public final class FederatedTrainingTaskDaoTest {
     }
 
     @After
-    public void cleanUp() throws Exception {
-        FederatedComputeDbHelper dbHelper =
-                FederatedComputeDbHelper.getInstanceForTest(mContext);
+    public void cleanUp() {
+        FederatedComputeDbHelper dbHelper = FederatedComputeDbHelper.getInstanceForTest(mContext);
         dbHelper.getWritableDatabase().close();
         dbHelper.getReadableDatabase().close();
         dbHelper.close();
@@ -96,6 +95,7 @@ public final class FederatedTrainingTaskDaoTest {
         mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task);
         FederatedTrainingTask task2 =
                 createDefaultFederatedTrainingTask().toBuilder()
+                        .jobId(456)
                         .populationName(POPULATION_NAME + "_2")
                         .build();
         mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task2);
@@ -122,6 +122,7 @@ public final class FederatedTrainingTaskDaoTest {
         mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task);
         FederatedTrainingTask task2 =
                 createDefaultFederatedTrainingTask().toBuilder()
+                        .jobId(456)
                         .populationName(POPULATION_NAME + "_2")
                         .build();
         mTrainingTaskDao.updateOrInsertFederatedTrainingTask(task2);
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java
new file mode 100644
index 00000000..0d1e410c
--- /dev/null
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/BackgroundKeyFetchJobServiceTest.java
@@ -0,0 +1,243 @@
+/*
+ * 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.encryption;
+
+import static com.android.federatedcompute.services.data.FederatedComputeEncryptionKey.KEY_TYPE_ENCRYPTION;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertFalse;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.app.job.JobInfo;
+import android.app.job.JobParameters;
+import android.app.job.JobScheduler;
+import android.content.Context;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import com.android.dx.mockito.inline.extended.ExtendedMockito;
+import com.android.federatedcompute.services.common.FederatedComputeExecutors;
+import com.android.federatedcompute.services.common.FederatedComputeJobInfo;
+import com.android.federatedcompute.services.common.FlagsFactory;
+import com.android.federatedcompute.services.common.MonotonicClock;
+import com.android.federatedcompute.services.common.PhFlagsTestUtil;
+import com.android.federatedcompute.services.data.FederatedComputeDbHelper;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKey;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKeyDao;
+import com.android.federatedcompute.services.http.HttpClient;
+
+import com.google.common.util.concurrent.FluentFuture;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+
+// TODO: add tests with Ph flags
+@RunWith(JUnit4.class)
+public class BackgroundKeyFetchJobServiceTest {
+
+    private BackgroundKeyFetchJobService mSpyService;
+
+    private MockitoSession mStaticMockSession;
+
+    private Context mContext;
+
+    private HttpClient mHttpClient;
+
+    public FederatedComputeEncryptionKeyDao mEncryptionDao;
+
+    public FederatedComputeEncryptionKeyManager mSpyKeyManager;
+
+    private TestInjector mInjector;
+
+    @Before
+    public void setUp() throws Exception {
+        PhFlagsTestUtil.setUpDeviceConfigPermissions();
+        PhFlagsTestUtil.disableGlobalKillSwitch();
+        MockitoAnnotations.initMocks(this);
+        mContext = ApplicationProvider.getApplicationContext();
+        mInjector = new TestInjector();
+        mEncryptionDao = FederatedComputeEncryptionKeyDao.getInstanceForTest(mContext);
+        mHttpClient = new HttpClient();
+        mSpyService = spy(new BackgroundKeyFetchJobService(new TestInjector()));
+        doReturn(mSpyService).when(mSpyService).getApplicationContext();
+        doNothing().when(mSpyService).jobFinished(any(), anyBoolean());
+        JobScheduler jobScheduler = mContext.getSystemService(JobScheduler.class);
+        jobScheduler.cancel(FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);
+        mSpyKeyManager =
+                spy(
+                        new FederatedComputeEncryptionKeyManager(
+                                MonotonicClock.getInstance(),
+                                mEncryptionDao,
+                                FlagsFactory.getFlags(),
+                                mHttpClient,
+                                MoreExecutors.newDirectExecutorService()));
+        mStaticMockSession =
+                ExtendedMockito.mockitoSession()
+                        .initMocks(this)
+                        .strictness(Strictness.LENIENT)
+                        .startMocking();
+    }
+
+    @After
+    public void tearDown() {
+        if (mStaticMockSession != null) {
+            mStaticMockSession.finishMocking();
+        }
+
+        FederatedComputeDbHelper dbHelper = FederatedComputeDbHelper.getInstanceForTest(mContext);
+        dbHelper.getWritableDatabase().close();
+        dbHelper.getReadableDatabase().close();
+        dbHelper.close();
+    }
+
+    @Test
+    public void testOnStartJob() {
+        FederatedComputeEncryptionKeyManager keyManager =
+                mInjector.getEncryptionKeyManager(mContext);
+        List<FederatedComputeEncryptionKey> emptyKeyList = List.of();
+        doReturn(FluentFuture.from(Futures.immediateFuture(emptyKeyList)))
+                .when(keyManager)
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true);
+
+        mSpyService.run(mock(JobParameters.class));
+
+        verify(mSpyService, times(1)).onStartJob(any());
+        verify(mSpyService, times(1)).jobFinished(any(), anyBoolean());
+    }
+
+    @Test
+    public void testOnStartJob_onFailure() {
+        FederatedComputeEncryptionKeyManager keyManager =
+                mInjector.getEncryptionKeyManager(mContext);
+        doReturn(
+                        FluentFuture.from(
+                                Futures.immediateFailedFuture(
+                                        new ExecutionException(
+                                                " Failed to fetch keys",
+                                                new IllegalStateException("http 404")))))
+                .when(keyManager)
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true);
+
+        mSpyService.run(mock(JobParameters.class));
+
+        verify(mSpyService, times(1)).onStartJob(any());
+        verify(mSpyService, times(1)).jobFinished(any(), anyBoolean());
+    }
+
+    @Test
+    public void testScheduleJob() {
+        final JobScheduler jobScheduler = mContext.getSystemService(JobScheduler.class);
+
+        assertThat(
+                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
+                                mContext, FlagsFactory.getFlags()))
+                .isEqualTo(true);
+
+        final JobInfo scheduledJob =
+                jobScheduler.getPendingJob(
+                        FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);
+
+        assertThat(scheduledJob.getId())
+                .isEqualTo(FederatedComputeJobInfo.ENCRYPTION_KEY_FETCH_JOB_ID);
+    }
+
+    @Test
+    public void testScheduleJob_notNeeded() {
+        assertThat(
+                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
+                                mContext, FlagsFactory.getFlags()))
+                .isEqualTo(true);
+
+        assertThat(
+                        BackgroundKeyFetchJobService.scheduleJobIfNeeded(
+                                mContext, FlagsFactory.getFlags()))
+                .isEqualTo(false);
+    }
+
+    @Test
+    public void testOnStopJob() {
+        assertFalse(mSpyService.onStopJob(mock(JobParameters.class)));
+    }
+
+    @Test
+    public void testOnStartJob_enableKillSwitch() {
+        PhFlagsTestUtil.enableGlobalKillSwitch();
+        FederatedComputeEncryptionKeyManager keyManager =
+                mInjector.getEncryptionKeyManager(mContext);
+        List<FederatedComputeEncryptionKey> emptyKeyList = List.of();
+        doReturn(FluentFuture.from(Futures.immediateFuture(emptyKeyList)))
+                .when(keyManager)
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true);
+
+        mSpyService.run(mock(JobParameters.class));
+
+        verify(mSpyService, times(1)).onStartJob(any());
+        verify(mSpyService, times(1)).jobFinished(any(), anyBoolean());
+        verify(keyManager, never()).fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION,
+                /* isScheduledJob= */ true);
+    }
+
+    @Test
+    public void testDefaultInjector() {
+        BackgroundKeyFetchJobService.Injector injector =
+                new BackgroundKeyFetchJobService.Injector();
+        assertThat(injector.getExecutor())
+                .isEqualTo(FederatedComputeExecutors.getBackgroundExecutor());
+        assertThat(injector.getEncryptionKeyManager(mContext))
+                .isEqualTo(FederatedComputeEncryptionKeyManager.getInstance(mContext));
+        assertThat(injector.getLightWeightExecutor())
+                .isEqualTo(FederatedComputeExecutors.getLightweightExecutor());
+    }
+
+    class TestInjector extends BackgroundKeyFetchJobService.Injector {
+        @Override
+        ListeningExecutorService getExecutor() {
+            return MoreExecutors.newDirectExecutorService();
+        }
+
+        @Override
+        ListeningExecutorService getLightWeightExecutor() {
+            return MoreExecutors.newDirectExecutorService();
+        }
+
+        @Override
+        FederatedComputeEncryptionKeyManager getEncryptionKeyManager(Context context) {
+            return mSpyKeyManager;
+        }
+    }
+}
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/FederatedComputeKeyFetchManagerTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/FederatedComputeKeyFetchManagerTest.java
new file mode 100644
index 00000000..5a1471e1
--- /dev/null
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/encryption/FederatedComputeKeyFetchManagerTest.java
@@ -0,0 +1,452 @@
+/*
+ * 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.federatedcompute.services.encryption;
+
+import static com.android.federatedcompute.services.data.FederatedComputeEncryptionKey.KEY_TYPE_ENCRYPTION;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.federatedcompute.services.common.Clock;
+import com.android.federatedcompute.services.common.FederatedComputeExecutors;
+import com.android.federatedcompute.services.common.Flags;
+import com.android.federatedcompute.services.common.MonotonicClock;
+import com.android.federatedcompute.services.data.FederatedComputeDbHelper;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKey;
+import com.android.federatedcompute.services.data.FederatedComputeEncryptionKeyDao;
+import com.android.federatedcompute.services.http.FederatedComputeHttpResponse;
+import com.android.federatedcompute.services.http.HttpClient;
+
+import com.google.common.util.concurrent.Futures;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.MockitoAnnotations;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+
+@RunWith(AndroidJUnit4.class)
+public class FederatedComputeKeyFetchManagerTest {
+
+    private static final Map<String, List<String>> SAMPLE_RESPONSE_HEADER =
+            Map.of(
+                    "Cache-Control", List.of("public,max-age=6000"),
+                    "Age", List.of("1"),
+                    "Content-Type", List.of("json"));
+
+    private static final String SAMPLE_RESPONSE_PAYLOAD =
+            """
+{ "keys": [{ "id": "0cc9b4c9-08bd", "key": "BQo+c1Tw6TaQ+VH/b+9PegZOjHuKAFkl8QdmS0IjRj8" """
+                    + "} ] }";
+
+    private FederatedComputeEncryptionKeyManager mFederatedComputeEncryptionKeyManager;
+
+    @Mock private HttpClient mMockHttpClient;
+
+    private FederatedComputeEncryptionKeyDao mEncryptionKeyDao;
+
+    private Context mContext;
+
+    private Clock mClock;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mContext = ApplicationProvider.getApplicationContext();
+        mClock = MonotonicClock.getInstance();
+        mEncryptionKeyDao = FederatedComputeEncryptionKeyDao.getInstanceForTest(mContext);
+        Flags mockFlags = Mockito.mock(Flags.class);
+        mFederatedComputeEncryptionKeyManager =
+                new FederatedComputeEncryptionKeyManager(
+                        mClock,
+                        mEncryptionKeyDao,
+                        mockFlags,
+                        mMockHttpClient,
+                        FederatedComputeExecutors.getBackgroundExecutor());
+        String overrideUrl = "https://real-coordinator/v1alpha/publicKeys";
+        doReturn(overrideUrl).when(mockFlags).getEncryptionKeyFetchUrl();
+    }
+
+    @After
+    public void teadDown() {
+        FederatedComputeDbHelper dbHelper = FederatedComputeDbHelper.getInstanceForTest(mContext);
+        dbHelper.getWritableDatabase().close();
+        dbHelper.getReadableDatabase().close();
+        dbHelper.close();
+    }
+
+    @Test
+    public void testGetTTL_fullInfo() {
+        Map<String, List<String>> headers = new HashMap<>();
+        headers.put("Cache-Control", List.of("public,max-age=3600"));
+        headers.put("Age", List.of("8"));
+
+        long ttl = FederatedComputeEncryptionKeyManager.getTTL(headers);
+
+        assertThat(ttl).isEqualTo(3600 - 8);
+    }
+
+    @Test
+    public void testGetTTL_noCache() {
+        Map<String, List<String>> headers = new HashMap<>();
+        headers.put("Age", List.of("8"));
+
+        long ttl = FederatedComputeEncryptionKeyManager.getTTL(headers);
+
+        assertThat(ttl).isEqualTo(0);
+    }
+
+    @Test
+    public void testGetTTL_noAge() {
+        Map<String, List<String>> headers = new HashMap<>();
+        headers.put("Cache-Control", List.of("public,max-age=3600"));
+
+        long ttl = FederatedComputeEncryptionKeyManager.getTTL(headers);
+
+        assertThat(ttl).isEqualTo(3600);
+    }
+
+    @Test
+    public void testGetTTL_empty() {
+        Map<String, List<String>> headers = Collections.EMPTY_MAP;
+
+        long ttl = FederatedComputeEncryptionKeyManager.getTTL(headers);
+
+        assertThat(ttl).isEqualTo(0);
+    }
+
+    @Test
+    public void testGetTTL_failedParse() {
+        Map<String, List<String>> headers = new HashMap<>();
+        headers.put("Cache-Control", List.of("public,max-age==3600"));
+        headers.put("Age", List.of("8"));
+
+        long ttl = FederatedComputeEncryptionKeyManager.getTTL(headers);
+
+        assertThat(ttl).isEqualTo(0);
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_scheduled_success() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        List<FederatedComputeEncryptionKey> keys =
+                mFederatedComputeEncryptionKeyManager
+                        .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true)
+                        .get();
+
+        assertThat(keys.size()).isGreaterThan(0);
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_nonScheduled_success() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        List<FederatedComputeEncryptionKey> keys =
+                mFederatedComputeEncryptionKeyManager
+                        .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ false)
+                        .get();
+
+        assertThat(keys.size()).isGreaterThan(0);
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_scheduled_throws() {
+        doReturn(
+                        Futures.immediateFailedFuture(
+                                new ExecutionException(
+                                        "fetchAndPersistActiveKeys keys failed.",
+                                        new IllegalStateException("http 404"))))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        assertThrows(
+                ExecutionException.class,
+                () ->
+                        mFederatedComputeEncryptionKeyManager
+                                .fetchAndPersistActiveKeys(
+                                        KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true)
+                                .get());
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_nonScheduled_throws() {
+        doReturn(
+                        Futures.immediateFailedFuture(
+                                new ExecutionException(
+                                        "fetchAndPersistActiveKeys keys failed.",
+                                        new IllegalStateException("http 404"))))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        assertThrows(
+                ExecutionException.class,
+                () ->
+                        mFederatedComputeEncryptionKeyManager
+                                .fetchAndPersistActiveKeys(
+                                        KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ false)
+                                .get());
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_scheduledNoDeletion() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        mFederatedComputeEncryptionKeyManager
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true)
+                .get();
+        List<FederatedComputeEncryptionKey> keys =
+                mEncryptionKeyDao.readFederatedComputeEncryptionKeysFromDatabase(
+                        ""
+                        /* selection= */ ,
+                        new String[0]
+                        /* selectionArgs= */ ,
+                        ""
+                        /* orderBy= */ ,
+                        -1
+                        /* count= */);
+
+        assertThat(keys.size()).isEqualTo(1);
+        assertThat(
+                        keys.stream()
+                                .map(FederatedComputeEncryptionKey::getKeyIdentifier)
+                                .collect(Collectors.toList()))
+                .containsAtLeastElementsIn(List.of("0cc9b4c9-08bd"));
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_nonScheduledNoDeletion() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        mFederatedComputeEncryptionKeyManager
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ false)
+                .get();
+        List<FederatedComputeEncryptionKey> keys =
+                mEncryptionKeyDao.readFederatedComputeEncryptionKeysFromDatabase(
+                        ""
+                        /* selection= */ ,
+                        new String[0]
+                        /* selectionArgs= */ ,
+                        ""
+                        /* orderBy= */ ,
+                        -1
+                        /* count= */);
+
+        assertThat(keys.size()).isEqualTo(1);
+        assertThat(
+                        keys.stream()
+                                .map(FederatedComputeEncryptionKey::getKeyIdentifier)
+                                .collect(Collectors.toList()))
+                .containsAtLeastElementsIn(List.of("0cc9b4c9-08bd"));
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_scheduledWithDeletion() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+        long currentTime = mClock.currentTimeMillis();
+        mEncryptionKeyDao.insertEncryptionKey(
+                new FederatedComputeEncryptionKey.Builder()
+                        .setKeyIdentifier("5161e286-63e5")
+                        .setPublicKey("YuOorP14obQLqASrvqbkNxyijjcAUIDx/xeMGZOyykc")
+                        .setKeyType(KEY_TYPE_ENCRYPTION)
+                        .setCreationTime(currentTime)
+                        .setExpiryTime(currentTime)
+                        .build());
+
+        mFederatedComputeEncryptionKeyManager
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ true)
+                .get();
+
+        List<FederatedComputeEncryptionKey> keys =
+                mEncryptionKeyDao.readFederatedComputeEncryptionKeysFromDatabase(
+                        ""
+                        /* selection= */ ,
+                        new String[0]
+                        /* selectionArgs= */ ,
+                        ""
+                        /* orderBy= */ ,
+                        -1
+                        /* count= */);
+
+        assertThat(keys.size()).isEqualTo(1);
+    }
+
+    @Test
+    public void testFetchAndPersistActiveKeys_nonScheduledWithDeletion() throws Exception {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+        long currentTime = mClock.currentTimeMillis();
+        mEncryptionKeyDao.insertEncryptionKey(
+                new FederatedComputeEncryptionKey.Builder()
+                        .setKeyIdentifier("5161e286-63e5")
+                        .setPublicKey("YuOorP14obQLqASrvqbkNxyijjcAUIDx/xeMGZOyykc")
+                        .setKeyType(KEY_TYPE_ENCRYPTION)
+                        .setCreationTime(currentTime)
+                        .setExpiryTime(currentTime)
+                        .build());
+
+        mFederatedComputeEncryptionKeyManager
+                .fetchAndPersistActiveKeys(KEY_TYPE_ENCRYPTION, /* isScheduledJob= */ false)
+                .get();
+
+        List<FederatedComputeEncryptionKey> keys =
+                mEncryptionKeyDao.readFederatedComputeEncryptionKeysFromDatabase(
+                        ""
+                        /* selection= */ ,
+                        new String[0]
+                        /* selectionArgs= */ ,
+                        ""
+                        /* orderBy= */ ,
+                        -1
+                        /* count= */);
+
+        assertThat(keys.size()).isEqualTo(2);
+
+        List<FederatedComputeEncryptionKey> activeKeys = mEncryptionKeyDao.getLatestExpiryNKeys(2);
+        assertThat(activeKeys.size()).isEqualTo(1);
+    }
+
+    @Test
+    public void testGetOrFetchActiveKeys_fetch() {
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        List<FederatedComputeEncryptionKey> keys =
+                mFederatedComputeEncryptionKeyManager.getOrFetchActiveKeys(
+                        KEY_TYPE_ENCRYPTION, /* keyCount= */ 2);
+
+        verify(mMockHttpClient, times(1)).performRequestAsyncWithRetry(any());
+        assertThat(keys.size()).isEqualTo(1);
+    }
+
+    @Test
+    public void testGetOrFetchActiveKeys_noFetch() {
+        long currentTime = mClock.currentTimeMillis();
+        mEncryptionKeyDao.insertEncryptionKey(
+                new FederatedComputeEncryptionKey.Builder()
+                        .setKeyIdentifier("5161e286-63e5")
+                        .setPublicKey("YuOorP14obQLqASrvqbkNxyijjcAUIDx/xeMGZOyykc")
+                        .setKeyType(KEY_TYPE_ENCRYPTION)
+                        .setCreationTime(currentTime)
+                        .setExpiryTime(currentTime + 5000L)
+                        .build());
+        doReturn(
+                        Futures.immediateFuture(
+                                new FederatedComputeHttpResponse.Builder()
+                                        .setHeaders(SAMPLE_RESPONSE_HEADER)
+                                        .setPayload(SAMPLE_RESPONSE_PAYLOAD.getBytes())
+                                        .setStatusCode(200)
+                                        .build()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        List<FederatedComputeEncryptionKey> keys =
+                mFederatedComputeEncryptionKeyManager.getOrFetchActiveKeys(
+                        KEY_TYPE_ENCRYPTION, /* keyCount= */ 2);
+
+        verify(mMockHttpClient, never()).performRequestAsyncWithRetry(any());
+        assertThat(keys.size()).isEqualTo(1);
+    }
+
+    @Test
+    public void testGetOrFetchActiveKeys_failure() {
+        doReturn(Futures.immediateFailedFuture(new InterruptedException()))
+                .when(mMockHttpClient)
+                .performRequestAsyncWithRetry(any());
+
+        List<FederatedComputeEncryptionKey> keys =
+                mFederatedComputeEncryptionKeyManager.getOrFetchActiveKeys(
+                        KEY_TYPE_ENCRYPTION, /* keyCount= */ 2);
+
+        assertThat(keys.size()).isEqualTo(0);
+        verify(mMockHttpClient, times(1)).performRequestAsyncWithRetry(any());
+    }
+}
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequestTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequestTest.java
index a504df65..fb5a339c 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequestTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpRequestTest.java
@@ -16,6 +16,9 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.http.HttpClientUtil.ACCEPT_ENCODING_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.GZIP_ENCODING_HDR;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertThrows;
@@ -39,11 +42,7 @@ public final class FederatedComputeHttpRequestTest {
                 IllegalArgumentException.class,
                 () ->
                         FederatedComputeHttpRequest.create(
-                                "http://invalid.com",
-                                HttpMethod.GET,
-                                new HashMap<>(),
-                                PAYLOAD,
-                                /* useCompression= */ false));
+                                "http://invalid.com", HttpMethod.GET, new HashMap<>(), PAYLOAD));
     }
 
     @Test
@@ -52,11 +51,7 @@ public final class FederatedComputeHttpRequestTest {
                 IllegalArgumentException.class,
                 () ->
                         FederatedComputeHttpRequest.create(
-                                "https://valid.com",
-                                HttpMethod.GET,
-                                new HashMap<>(),
-                                PAYLOAD,
-                                /* useCompression= */ false));
+                                "https://valid.com", HttpMethod.GET, new HashMap<>(), PAYLOAD));
     }
 
     @Test
@@ -67,11 +62,7 @@ public final class FederatedComputeHttpRequestTest {
                 IllegalArgumentException.class,
                 () ->
                         FederatedComputeHttpRequest.create(
-                                "https://valid.com",
-                                HttpMethod.POST,
-                                headers,
-                                PAYLOAD,
-                                /* useCompression= */ false));
+                                "https://valid.com", HttpMethod.POST, headers, PAYLOAD));
     }
 
     @Test
@@ -79,11 +70,7 @@ public final class FederatedComputeHttpRequestTest {
         String expectedUri = "https://valid.com";
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.GET,
-                        new HashMap<>(),
-                        HttpClientUtil.EMPTY_BODY,
-                        /* useCompression= */ false);
+                        expectedUri, HttpMethod.GET, new HashMap<>(), HttpClientUtil.EMPTY_BODY);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertThat(request.getHttpMethod()).isEqualTo(HttpMethod.GET);
@@ -99,27 +86,19 @@ public final class FederatedComputeHttpRequestTest {
 
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.GET,
-                        expectedHeaders,
-                        HttpClientUtil.EMPTY_BODY,
-                        /* useCompression= */ false);
+                        expectedUri, HttpMethod.GET, expectedHeaders, HttpClientUtil.EMPTY_BODY);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertThat(request.getExtraHeaders()).isEqualTo(expectedHeaders);
     }
 
     @Test
-    public void createPostRequestWithoutBody_valid() throws Exception {
+    public void createPostRequestWithoutBody_valid() {
         String expectedUri = "https://valid.com";
 
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.POST,
-                        new HashMap<>(),
-                        HttpClientUtil.EMPTY_BODY,
-                        /* useCompression= */ false);
+                        expectedUri, HttpMethod.POST, new HashMap<>(), HttpClientUtil.EMPTY_BODY);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertTrue(request.getExtraHeaders().isEmpty());
@@ -128,16 +107,12 @@ public final class FederatedComputeHttpRequestTest {
     }
 
     @Test
-    public void createPostRequestWithBody_valid() throws Exception {
+    public void createPostRequestWithBody_valid() {
         String expectedUri = "https://valid.com";
 
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.POST,
-                        new HashMap<>(),
-                        PAYLOAD,
-                        /* useCompression= */ false);
+                        expectedUri, HttpMethod.POST, new HashMap<>(), PAYLOAD);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertThat(request.getHttpMethod()).isEqualTo(HttpMethod.POST);
@@ -145,18 +120,14 @@ public final class FederatedComputeHttpRequestTest {
     }
 
     @Test
-    public void createPostRequestWithBodyHeader_valid() throws Exception {
+    public void createPostRequestWithBodyHeader_valid() {
         String expectedUri = "https://valid.com";
         HashMap<String, String> expectedHeaders = new HashMap<>();
         expectedHeaders.put("Foo", "Bar");
 
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.POST,
-                        expectedHeaders,
-                        PAYLOAD,
-                        /* useCompression= */ false);
+                        expectedUri, HttpMethod.POST, expectedHeaders, PAYLOAD);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertThat(request.getHttpMethod()).isEqualTo(HttpMethod.POST);
@@ -165,23 +136,19 @@ public final class FederatedComputeHttpRequestTest {
     }
 
     @Test
-    public void createPostRequestWithCompression_valid() throws Exception {
+    public void createGetRequestWithAcceptCompression_valid() {
         String expectedUri = "https://valid.com";
-
+        HashMap<String, String> headerList = new HashMap<>();
+        headerList.put(ACCEPT_ENCODING_HDR, GZIP_ENCODING_HDR);
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
-                        expectedUri,
-                        HttpMethod.POST,
-                        new HashMap<>(),
-                        PAYLOAD,
-                        /* useCompression= */ true);
+                        expectedUri, HttpMethod.POST, headerList, PAYLOAD);
 
         assertThat(request.getUri()).isEqualTo(expectedUri);
         assertThat(request.getHttpMethod()).isEqualTo(HttpMethod.POST);
-        assertThat(request.getBody()).isEqualTo(HttpClientUtil.compressWithGzip(PAYLOAD));
         HashMap<String, String> expectedHeaders = new HashMap<>();
-        expectedHeaders.put(HttpClientUtil.CONTENT_ENCODING_HDR, HttpClientUtil.GZIP_ENCODING_HDR);
-        expectedHeaders.put(HttpClientUtil.CONTENT_LENGTH_HDR, String.valueOf(42));
+        expectedHeaders.put(HttpClientUtil.CONTENT_LENGTH_HDR, String.valueOf(22));
+        expectedHeaders.put(ACCEPT_ENCODING_HDR, GZIP_ENCODING_HDR);
         assertThat(request.getExtraHeaders()).isEqualTo(expectedHeaders);
     }
 }
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponseTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponseTest.java
index 2497b19d..d9936631 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponseTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/FederatedComputeHttpResponseTest.java
@@ -16,6 +16,8 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.http.HttpClientUtil.OCTET_STREAM;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertThrows;
@@ -29,6 +31,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -58,7 +61,7 @@ public final class FederatedComputeHttpResponseTest {
     }
 
     @Test
-    public void testBuildWithMinimalRequiredValues() throws Exception {
+    public void testBuildWithMinimalRequiredValues() {
         final int responseCode = 200;
         FederatedComputeHttpResponse response =
                 new FederatedComputeHttpResponse.Builder().setStatusCode(responseCode).build();
@@ -75,4 +78,20 @@ public final class FederatedComputeHttpResponseTest {
                                 .setPayload("payload".getBytes(UTF_8))
                                 .build());
     }
+
+    @Test
+    public void testGetBody_success() {
+        final byte[] uncompressedBody = "payload".getBytes(UTF_8);
+        Map<String, List<String>> expectedHeaders = new HashMap<>();
+        expectedHeaders.put(HttpClientUtil.CONTENT_TYPE_HDR, ImmutableList.of(OCTET_STREAM));
+
+        FederatedComputeHttpResponse response =
+                new FederatedComputeHttpResponse.Builder()
+                        .setStatusCode(200)
+                        .setPayload(uncompressedBody)
+                        .setHeaders(expectedHeaders)
+                        .build();
+
+        assertThat(response.getPayload()).isEqualTo(uncompressedBody);
+    }
 }
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientTest.java
index 2674f2d3..c616ec21 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientTest.java
@@ -20,8 +20,11 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -34,6 +37,7 @@ import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
@@ -47,9 +51,16 @@ import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 @RunWith(JUnit4.class)
 public final class HttpClientTest {
+    public static final FederatedComputeHttpRequest DEFAULT_GET_REQUEST =
+            FederatedComputeHttpRequest.create(
+                    "https://google.com",
+                    HttpMethod.GET,
+                    new HashMap<>(),
+                    HttpClientUtil.EMPTY_BODY);
     @Spy private HttpClient mHttpClient = new HttpClient();
     @Rule public MockitoRule rule = MockitoJUnit.rule();
     @Mock private HttpURLConnection mMockHttpURLConnection;
@@ -61,22 +72,14 @@ public final class HttpClientTest {
                         "https://google.com",
                         HttpMethod.POST,
                         new HashMap<>(),
-                        HttpClientUtil.EMPTY_BODY,
-                        false);
+                        HttpClientUtil.EMPTY_BODY);
         doThrow(new IOException()).when(mHttpClient).setup(ArgumentMatchers.any());
 
-        assertThrows(IllegalArgumentException.class, () -> mHttpClient.performRequest(request));
+        assertThrows(IOException.class, () -> mHttpClient.performRequest(request));
     }
 
     @Test
     public void testPerformGetRequestSuccess() throws Exception {
-        FederatedComputeHttpRequest request =
-                FederatedComputeHttpRequest.create(
-                        "https://google.com",
-                        HttpMethod.GET,
-                        new HashMap<>(),
-                        HttpClientUtil.EMPTY_BODY,
-                        false);
         String successMessage = "Success!";
         InputStream mockStream = new ByteArrayInputStream(successMessage.getBytes(UTF_8));
         Map<String, List<String>> mockHeaders = new HashMap<>();
@@ -88,7 +91,7 @@ public final class HttpClientTest {
         when(mMockHttpURLConnection.getContentLengthLong())
                 .thenReturn((long) successMessage.length());
 
-        FederatedComputeHttpResponse response = mHttpClient.performRequest(request);
+        FederatedComputeHttpResponse response = mHttpClient.performRequest(DEFAULT_GET_REQUEST);
 
         assertThat(response.getStatusCode()).isEqualTo(200);
         assertThat(response.getHeaders()).isEqualTo(mockHeaders);
@@ -98,13 +101,6 @@ public final class HttpClientTest {
     @Test
     public void testPerformGetRequestFails() throws Exception {
         String failureMessage = "FAIL!";
-        FederatedComputeHttpRequest request =
-                FederatedComputeHttpRequest.create(
-                        "https://google.com",
-                        HttpMethod.GET,
-                        new HashMap<>(),
-                        HttpClientUtil.EMPTY_BODY,
-                        false);
         InputStream mockStream = new ByteArrayInputStream(failureMessage.getBytes(UTF_8));
         when(mMockHttpURLConnection.getErrorStream()).thenReturn(mockStream);
         when(mMockHttpURLConnection.getResponseCode()).thenReturn(503);
@@ -113,22 +109,78 @@ public final class HttpClientTest {
         when(mMockHttpURLConnection.getContentLengthLong())
                 .thenReturn((long) failureMessage.length());
 
-        FederatedComputeHttpResponse response = mHttpClient.performRequest(request);
+        FederatedComputeHttpResponse response = mHttpClient.performRequest(DEFAULT_GET_REQUEST);
+
+        assertThat(response.getStatusCode()).isEqualTo(503);
+        assertTrue(response.getHeaders().isEmpty());
+        assertThat(response.getPayload()).isEqualTo(failureMessage.getBytes(UTF_8));
+    }
+
+    @Test
+    public void testPerformGetRequestFailsWithRetry() throws Exception {
+        String failureMessage = "FAIL!";
+        when(mMockHttpURLConnection.getErrorStream())
+                .then(invocation -> new ByteArrayInputStream(failureMessage.getBytes(UTF_8)));
+        when(mMockHttpURLConnection.getResponseCode()).thenReturn(503);
+        when(mMockHttpURLConnection.getHeaderFields()).thenReturn(new HashMap<>());
+        when(mMockHttpURLConnection.getContentLengthLong())
+                .thenReturn((long) failureMessage.length());
+        doReturn(mMockHttpURLConnection).when(mHttpClient).setup(ArgumentMatchers.any());
 
+        FederatedComputeHttpResponse response =
+                mHttpClient.performRequestWithRetry(DEFAULT_GET_REQUEST);
+
+        verify(mHttpClient, times(3)).performRequest(DEFAULT_GET_REQUEST);
         assertThat(response.getStatusCode()).isEqualTo(503);
         assertTrue(response.getHeaders().isEmpty());
         assertThat(response.getPayload()).isEqualTo(failureMessage.getBytes(UTF_8));
     }
 
     @Test
+    public void testPerformGetRequestSuccessWithRetry() throws Exception {
+        String failureMessage = "FAIL!";
+        InputStream mockStream = new ByteArrayInputStream(failureMessage.getBytes(UTF_8));
+        when(mMockHttpURLConnection.getErrorStream()).thenReturn(mockStream);
+        when(mMockHttpURLConnection.getResponseCode()).thenReturn(503);
+        when(mMockHttpURLConnection.getHeaderFields()).thenReturn(new HashMap<>());
+        HttpURLConnection mockSuccessfulHttpURLConnection = Mockito.mock(HttpURLConnection.class);
+        Map<String, List<String>> mockHeaders = new HashMap<>();
+        mockHeaders.put("Header1", Arrays.asList("Value1"));
+        when(mockSuccessfulHttpURLConnection.getOutputStream())
+                .thenReturn(new ByteArrayOutputStream());
+        when(mockSuccessfulHttpURLConnection.getResponseCode()).thenReturn(200);
+        when(mockSuccessfulHttpURLConnection.getHeaderFields()).thenReturn(mockHeaders);
+        final AtomicInteger countCall = new AtomicInteger();
+        doAnswer(
+                        invocation -> {
+                            int count = countCall.incrementAndGet();
+                            if (count < 3) {
+                                return mMockHttpURLConnection;
+                            } else {
+                                return mockSuccessfulHttpURLConnection;
+                            }
+                        })
+                .when(mHttpClient)
+                .setup(ArgumentMatchers.any());
+        when(mMockHttpURLConnection.getContentLengthLong())
+                .thenReturn((long) failureMessage.length());
+
+        FederatedComputeHttpResponse response =
+                mHttpClient.performRequestWithRetry(DEFAULT_GET_REQUEST);
+
+        verify(mHttpClient, times(3)).performRequest(DEFAULT_GET_REQUEST);
+        assertThat(response.getStatusCode()).isEqualTo(200);
+        assertThat(response.getHeaders()).isEqualTo(mockHeaders);
+    }
+
+    @Test
     public void testPerformPostRequestSuccess() throws Exception {
         FederatedComputeHttpRequest request =
                 FederatedComputeHttpRequest.create(
                         "https://google.com",
                         HttpMethod.POST,
                         new HashMap<>(),
-                        "payload".getBytes(UTF_8),
-                        false);
+                        "payload".getBytes(UTF_8));
         Map<String, List<String>> mockHeaders = new HashMap<>();
         mockHeaders.put("Header1", Arrays.asList("Value1"));
         when(mMockHttpURLConnection.getOutputStream()).thenReturn(new ByteArrayOutputStream());
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientUtilTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientUtilTest.java
new file mode 100644
index 00000000..1f16aecc
--- /dev/null
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpClientUtilTest.java
@@ -0,0 +1,54 @@
+/*
+ * 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.http;
+
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.net.Uri;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import org.junit.Test;
+
+import java.io.ByteArrayOutputStream;
+import java.io.InputStream;
+
+public class HttpClientUtilTest {
+    @Test
+    public void compress_uncompress_success() throws Exception {
+        String testUriPrefix =
+                "android.resource://com.android.ondevicepersonalization.federatedcomputetests/raw/";
+        Uri checkpointUri = Uri.parse(testUriPrefix + "federation_test_checkpoint_client");
+        Context context = ApplicationProvider.getApplicationContext();
+        InputStream in = context.getContentResolver().openInputStream(checkpointUri);
+        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+        byte[] buf = new byte[4096];
+        int bytesRead;
+        while ((bytesRead = in.read(buf)) != -1) {
+            outputStream.write(buf, 0, bytesRead);
+        }
+        byte[] dataBeforeCompress = outputStream.toByteArray();
+
+        byte[] dataAfterCompress = HttpClientUtil.compressWithGzip(dataBeforeCompress);
+        assertThat(dataAfterCompress.length).isLessThan(dataBeforeCompress.length);
+
+        byte[] unzipData = HttpClientUtil.uncompressWithGzip(dataAfterCompress);
+        assertThat(unzipData).isEqualTo(dataBeforeCompress);
+    }
+}
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpFederatedProtocolTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpFederatedProtocolTest.java
index fb1d95c4..e483cf22 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpFederatedProtocolTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/HttpFederatedProtocolTest.java
@@ -16,31 +16,46 @@
 
 package com.android.federatedcompute.services.http;
 
+import static com.android.federatedcompute.services.http.HttpClientUtil.ACCEPT_ENCODING_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_ENCODING_HDR;
 import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_LENGTH_HDR;
 import static com.android.federatedcompute.services.http.HttpClientUtil.CONTENT_TYPE_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.GZIP_ENCODING_HDR;
+import static com.android.federatedcompute.services.http.HttpClientUtil.ODP_IDEMPOTENCY_KEY;
 import static com.android.federatedcompute.services.http.HttpClientUtil.PROTOBUF_CONTENT_TYPE;
+import static com.android.federatedcompute.services.http.HttpClientUtil.compressWithGzip;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.nio.file.StandardCopyOption.REPLACE_EXISTING;
+
+import android.content.Context;
+import android.net.Uri;
+
+import androidx.test.core.app.ApplicationProvider;
 
 import com.android.federatedcompute.services.http.HttpClientUtil.HttpMethod;
 import com.android.federatedcompute.services.testutils.TrainingTestUtil;
 import com.android.federatedcompute.services.training.util.ComputationResult;
 
-import com.google.common.io.Files;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.intelligence.fcp.client.FLRunnerResult;
 import com.google.intelligence.fcp.client.FLRunnerResult.ContributionResult;
-import com.google.internal.federated.plan.ClientOnlyPlan;
 import com.google.internal.federatedcompute.v1.ClientVersion;
 import com.google.internal.federatedcompute.v1.RejectionInfo;
 import com.google.internal.federatedcompute.v1.Resource;
+import com.google.internal.federatedcompute.v1.ResourceCapabilities;
+import com.google.internal.federatedcompute.v1.ResourceCompressionFormat;
 import com.google.ondevicepersonalization.federatedcompute.proto.CreateTaskAssignmentRequest;
 import com.google.ondevicepersonalization.federatedcompute.proto.CreateTaskAssignmentResponse;
 import com.google.ondevicepersonalization.federatedcompute.proto.ReportResultRequest;
@@ -53,7 +68,7 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.junit.runners.Parameterized;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
@@ -61,11 +76,14 @@ import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
 import java.io.File;
+import java.io.InputStream;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 
-@RunWith(JUnit4.class)
+@RunWith(Parameterized.class)
 public final class HttpFederatedProtocolTest {
     private static final String TASK_ASSIGNMENT_TARGET_URI = "https://test-server.com/";
     private static final String PLAN_URI = "https://fake.uri/plan";
@@ -85,18 +103,21 @@ public final class HttpFederatedProtocolTest {
     private static final String OCTET_STREAM = "application/octet-stream";
     private static final FLRunnerResult FL_RUNNER_SUCCESS_RESULT =
             FLRunnerResult.newBuilder().setContributionResult(ContributionResult.SUCCESS).build();
+
     private static final FLRunnerResult FL_RUNNER_FAIL_RESULT =
             FLRunnerResult.newBuilder().setContributionResult(ContributionResult.FAIL).build();
-    private static final FederatedComputeHttpResponse CHECKPOINT_HTTP_RESPONSE =
-            new FederatedComputeHttpResponse.Builder()
-                    .setStatusCode(200)
-                    .setPayload(CHECKPOINT)
-                    .build();
-
-    private static final CreateTaskAssignmentRequest START_TASK_ASSIGNMENT_REQUEST =
-            CreateTaskAssignmentRequest.newBuilder()
-                    .setClientVersion(ClientVersion.newBuilder().setVersionCode(CLIENT_VERSION))
-                    .build();
+    private static final CreateTaskAssignmentRequest
+            START_TASK_ASSIGNMENT_REQUEST_WITH_COMPRESSION =
+                    CreateTaskAssignmentRequest.newBuilder()
+                            .setClientVersion(
+                                    ClientVersion.newBuilder().setVersionCode(CLIENT_VERSION))
+                            .setResourceCapabilities(
+                                    ResourceCapabilities.newBuilder()
+                                            .addSupportedCompressionFormats(
+                                                    ResourceCompressionFormat
+                                                            .RESOURCE_COMPRESSION_FORMAT_GZIP)
+                                            .build())
+                            .build();
 
     private static final FederatedComputeHttpResponse SUCCESS_EMPTY_HTTP_RESPONSE =
             new FederatedComputeHttpResponse.Builder().setStatusCode(200).build();
@@ -105,6 +126,15 @@ public final class HttpFederatedProtocolTest {
 
     @Rule public MockitoRule rule = MockitoJUnit.rule();
     @Mock private HttpClient mMockHttpClient;
+
+    @Parameterized.Parameter(0)
+    public boolean mSupportCompression;
+
+    @Parameterized.Parameters
+    public static Collection<Object[]> data() {
+        return Arrays.asList(new Object[][] {{false}, {true}});
+    }
+
     private HttpFederatedProtocol mHttpFederatedProtocol;
 
     @Before
@@ -128,21 +158,39 @@ public final class HttpFederatedProtocolTest {
         // Verify task assignment request.
         FederatedComputeHttpRequest actualStartTaskAssignmentRequest = actualHttpRequests.get(0);
         assertThat(actualStartTaskAssignmentRequest.getUri()).isEqualTo(START_TASK_ASSIGNMENT_URI);
-        assertThat(actualStartTaskAssignmentRequest.getBody())
-                .isEqualTo(START_TASK_ASSIGNMENT_REQUEST.toByteArray());
+
         assertThat(actualStartTaskAssignmentRequest.getHttpMethod()).isEqualTo(HttpMethod.POST);
         HashMap<String, String> expectedHeaders = new HashMap<>();
-        expectedHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(18));
+        assertThat(actualStartTaskAssignmentRequest.getBody())
+                .isEqualTo(START_TASK_ASSIGNMENT_REQUEST_WITH_COMPRESSION.toByteArray());
+        expectedHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(23));
         expectedHeaders.put(CONTENT_TYPE_HDR, PROTOBUF_CONTENT_TYPE);
         assertThat(actualStartTaskAssignmentRequest.getExtraHeaders())
-                .containsExactlyEntriesIn(expectedHeaders);
+                .containsAtLeastEntriesIn(expectedHeaders);
+        assertThat(actualStartTaskAssignmentRequest.getExtraHeaders()).hasSize(3);
+        String idempotencyKey =
+                actualStartTaskAssignmentRequest.getExtraHeaders().get(ODP_IDEMPOTENCY_KEY);
+        assertNotNull(idempotencyKey);
+        String timestamp = idempotencyKey.split(" - ")[0];
+        assertThat(Long.parseLong(timestamp)).isLessThan(System.currentTimeMillis());
+
+        // Verify fetch resource request.
+        FederatedComputeHttpRequest actualFetchResourceRequest = actualHttpRequests.get(1);
+        ImmutableSet<String> resourceUris = ImmutableSet.of(PLAN_URI, CHECKPOINT_URI);
+        assertTrue(resourceUris.contains(actualFetchResourceRequest.getUri()));
+        expectedHeaders = new HashMap<>();
+        if (mSupportCompression) {
+            expectedHeaders.put(ACCEPT_ENCODING_HDR, GZIP_ENCODING_HDR);
+        }
+        assertThat(actualFetchResourceRequest.getExtraHeaders()).isEqualTo(expectedHeaders);
     }
 
     @Test
-    public void testCreateTaskAssignmentFailed() throws Exception {
+    public void testCreateTaskAssignmentFailed() {
         FederatedComputeHttpResponse httpResponse =
                 new FederatedComputeHttpResponse.Builder().setStatusCode(404).build();
-        when(mMockHttpClient.performRequestAsync(any())).thenReturn(immediateFuture(httpResponse));
+        when(mMockHttpClient.performRequestAsyncWithRetry(any()))
+                .thenReturn(immediateFuture(httpResponse));
 
         ExecutionException exception =
                 assertThrows(
@@ -166,15 +214,13 @@ public final class HttpFederatedProtocolTest {
                         .setStatusCode(200)
                         .setPayload(createTaskAssignmentResponse.toByteArray())
                         .build();
-        when(mMockHttpClient.performRequestAsync(any())).thenReturn(immediateFuture(httpResponse));
+        when(mMockHttpClient.performRequestAsyncWithRetry(any()))
+                .thenReturn(immediateFuture(httpResponse));
 
-        ExecutionException exception =
-                assertThrows(
-                        ExecutionException.class,
-                        () -> mHttpFederatedProtocol.issueCheckin().get());
+        CheckinResult checkinResult = mHttpFederatedProtocol.issueCheckin().get();
 
-        assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class);
-        assertThat(exception.getCause()).hasMessageThat().isEqualTo("Device rejected by server.");
+        assertThat(checkinResult.getRejectionInfo()).isNotNull();
+        assertThat(checkinResult.getRejectionInfo()).isEqualTo(RejectionInfo.getDefaultInstance());
     }
 
     @Test
@@ -186,7 +232,7 @@ public final class HttpFederatedProtocolTest {
         setUpHttpFederatedProtocol(
                 createStartTaskAssignmentHttpResponse(),
                 planHttpResponse,
-                CHECKPOINT_HTTP_RESPONSE,
+                checkpointHttpResponse(),
                 /** reportResultHttpResponse= */
                 null,
                 /** uploadResultHttpResponse= */
@@ -276,7 +322,12 @@ public final class HttpFederatedProtocolTest {
         assertThat(actualDataUploadRequest.getUri()).isEqualTo(UPLOAD_LOCATION_URI);
         assertThat(acutalReportResultRequest.getHttpMethod()).isEqualTo(HttpMethod.PUT);
         expectedHeaders = new HashMap<>();
-        expectedHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(17));
+        if (mSupportCompression) {
+            expectedHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(339));
+            expectedHeaders.put(CONTENT_ENCODING_HDR, GZIP_ENCODING_HDR);
+        } else {
+            expectedHeaders.put(CONTENT_LENGTH_HDR, String.valueOf(31846));
+        }
         expectedHeaders.put(CONTENT_TYPE_HDR, OCTET_STREAM);
         assertThat(actualDataUploadRequest.getExtraHeaders()).isEqualTo(expectedHeaders);
     }
@@ -291,7 +342,7 @@ public final class HttpFederatedProtocolTest {
         setUpHttpFederatedProtocol(
                 createStartTaskAssignmentHttpResponse(),
                 createPlanHttpResponse(),
-                CHECKPOINT_HTTP_RESPONSE,
+                checkpointHttpResponse(),
                 reportResultHttpResponse,
                 null);
 
@@ -315,7 +366,7 @@ public final class HttpFederatedProtocolTest {
         setUpHttpFederatedProtocol(
                 createStartTaskAssignmentHttpResponse(),
                 createPlanHttpResponse(),
-                CHECKPOINT_HTTP_RESPONSE,
+                checkpointHttpResponse(),
                 createReportResultHttpResponse(),
                 uploadResultHttpResponse);
 
@@ -330,41 +381,50 @@ public final class HttpFederatedProtocolTest {
     }
 
     private String createOutputCheckpointFile() throws Exception {
+        String testUriPrefix =
+                "android.resource://com.android.ondevicepersonalization.federatedcomputetests/raw/";
         File outputCheckpointFile = File.createTempFile("output", ".ckp");
-        Files.write("output checkpoint".getBytes(), outputCheckpointFile);
+        Context context = ApplicationProvider.getApplicationContext();
+        Uri checkpointUri = Uri.parse(testUriPrefix + "federation_test_checkpoint_client");
+        InputStream in = context.getContentResolver().openInputStream(checkpointUri);
+        java.nio.file.Files.copy(in, outputCheckpointFile.toPath(), REPLACE_EXISTING);
+        in.close();
         outputCheckpointFile.deleteOnExit();
         return outputCheckpointFile.getAbsolutePath();
     }
 
     private FederatedComputeHttpResponse createPlanHttpResponse() {
-        ClientOnlyPlan clientOnlyPlan = TrainingTestUtil.createFederatedAnalyticClientPlan();
+        byte[] clientOnlyPlan = TrainingTestUtil.createFederatedAnalyticClientPlan().toByteArray();
         return new FederatedComputeHttpResponse.Builder()
                 .setStatusCode(200)
-                .setPayload(clientOnlyPlan.toByteArray())
+                .setHeaders(mSupportCompression ? compressionHeaderList() : new HashMap<>())
+                .setPayload(mSupportCompression ? compressWithGzip(clientOnlyPlan) : clientOnlyPlan)
                 .build();
     }
 
-    private void setUpHttpFederatedProtocol() throws Exception {
-        FederatedComputeHttpResponse checkpointHttpResponse =
-                new FederatedComputeHttpResponse.Builder()
-                        .setStatusCode(200)
-                        .setPayload(CHECKPOINT)
-                        .build();
+    private void setUpHttpFederatedProtocol() {
         setUpHttpFederatedProtocol(
                 createStartTaskAssignmentHttpResponse(),
                 createPlanHttpResponse(),
-                checkpointHttpResponse,
+                checkpointHttpResponse(),
                 createReportResultHttpResponse(),
                 SUCCESS_EMPTY_HTTP_RESPONSE);
     }
 
+    private FederatedComputeHttpResponse checkpointHttpResponse() {
+        return new FederatedComputeHttpResponse.Builder()
+                .setStatusCode(200)
+                .setPayload(mSupportCompression ? compressWithGzip(CHECKPOINT) : CHECKPOINT)
+                .setHeaders(mSupportCompression ? compressionHeaderList() : new HashMap<>())
+                .build();
+    }
+
     private void setUpHttpFederatedProtocol(
             FederatedComputeHttpResponse createTaskAssignmentResponse,
             FederatedComputeHttpResponse planHttpResponse,
             FederatedComputeHttpResponse checkpointHttpResponse,
             FederatedComputeHttpResponse reportResultHttpResponse,
-            FederatedComputeHttpResponse uploadResultHttpResponse)
-            throws Exception {
+            FederatedComputeHttpResponse uploadResultHttpResponse) {
         doAnswer(
                         invocation -> {
                             FederatedComputeHttpRequest httpRequest = invocation.getArgument(0);
@@ -383,13 +443,25 @@ public final class HttpFederatedProtocolTest {
                             return immediateFuture(SUCCESS_EMPTY_HTTP_RESPONSE);
                         })
                 .when(mMockHttpClient)
-                .performRequestAsync(mHttpRequestCaptor.capture());
+                .performRequestAsyncWithRetry(mHttpRequestCaptor.capture());
+    }
+
+    private HashMap<String, List<String>> compressionHeaderList() {
+        HashMap<String, List<String>> headerList = new HashMap<>();
+        headerList.put(CONTENT_ENCODING_HDR, ImmutableList.of(GZIP_ENCODING_HDR));
+        headerList.put(CONTENT_TYPE_HDR, ImmutableList.of(OCTET_STREAM));
+        return headerList;
     }
 
-    private FederatedComputeHttpResponse createReportResultHttpResponse() throws Exception {
+    private FederatedComputeHttpResponse createReportResultHttpResponse() {
         UploadInstruction.Builder uploadInstruction =
                 UploadInstruction.newBuilder().setUploadLocation(UPLOAD_LOCATION_URI);
         uploadInstruction.putExtraRequestHeaders(CONTENT_TYPE_HDR, OCTET_STREAM);
+        if (mSupportCompression) {
+            uploadInstruction.putExtraRequestHeaders(CONTENT_ENCODING_HDR, GZIP_ENCODING_HDR);
+            uploadInstruction.setCompressionFormat(
+                    ResourceCompressionFormat.RESOURCE_COMPRESSION_FORMAT_GZIP);
+        }
         ReportResultResponse reportResultResponse =
                 ReportResultResponse.newBuilder()
                         .setUploadInstruction(uploadInstruction.build())
@@ -400,11 +472,27 @@ public final class HttpFederatedProtocolTest {
                 .build();
     }
 
-    private FederatedComputeHttpResponse createStartTaskAssignmentHttpResponse() throws Exception {
+    private FederatedComputeHttpResponse createStartTaskAssignmentHttpResponse() {
         CreateTaskAssignmentResponse createTaskAssignmentResponse =
                 createCreateTaskAssignmentResponse(
-                        Resource.newBuilder().setUri(PLAN_URI).build(),
-                        Resource.newBuilder().setUri(CHECKPOINT_URI).build());
+                        Resource.newBuilder()
+                                .setUri(PLAN_URI)
+                                .setCompressionFormat(
+                                        mSupportCompression
+                                                ? ResourceCompressionFormat
+                                                        .RESOURCE_COMPRESSION_FORMAT_GZIP
+                                                : ResourceCompressionFormat
+                                                        .RESOURCE_COMPRESSION_FORMAT_UNSPECIFIED)
+                                .build(),
+                        Resource.newBuilder()
+                                .setUri(CHECKPOINT_URI)
+                                .setCompressionFormat(
+                                        mSupportCompression
+                                                ? ResourceCompressionFormat
+                                                        .RESOURCE_COMPRESSION_FORMAT_GZIP
+                                                : ResourceCompressionFormat
+                                                        .RESOURCE_COMPRESSION_FORMAT_UNSPECIFIED)
+                                .build());
 
         return new FederatedComputeHttpResponse.Builder()
                 .setStatusCode(200)
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/ProtocolRequestCreatorTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/ProtocolRequestCreatorTest.java
index 1fa86e40..afcc97c9 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/ProtocolRequestCreatorTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/http/ProtocolRequestCreatorTest.java
@@ -43,7 +43,7 @@ public final class ProtocolRequestCreatorTest {
     @Test
     public void testCreateProtobufEncodedRequest() {
         ProtocolRequestCreator requestCreator =
-                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>(), false);
+                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>());
 
         FederatedComputeHttpRequest request =
                 requestCreator.createProtoRequest(
@@ -63,9 +63,7 @@ public final class ProtocolRequestCreatorTest {
         IllegalArgumentException exception =
                 assertThrows(
                         IllegalArgumentException.class,
-                        () ->
-                                ProtocolRequestCreator.create(
-                                        ForwardingInfo.getDefaultInstance(), false));
+                        () -> ProtocolRequestCreator.create(ForwardingInfo.getDefaultInstance()));
 
         assertThat(exception)
                 .hasMessageThat()
@@ -75,7 +73,7 @@ public final class ProtocolRequestCreatorTest {
     @Test
     public void testCreateProtocolRequestInvalidSuffix() {
         ProtocolRequestCreator requestCreator =
-                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>(), false);
+                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>());
 
         IllegalArgumentException exception =
                 assertThrows(
@@ -93,8 +91,7 @@ public final class ProtocolRequestCreatorTest {
     public void testCreateProtocolRequestWithForwardingInfo() {
         ForwardingInfo forwardingInfo =
                 ForwardingInfo.newBuilder().setTargetUriPrefix(AGGREGATION_TARGET_URI).build();
-        ProtocolRequestCreator requestCreator =
-                ProtocolRequestCreator.create(forwardingInfo, false);
+        ProtocolRequestCreator requestCreator = ProtocolRequestCreator.create(forwardingInfo);
 
         FederatedComputeHttpRequest request =
                 requestCreator.createProtoRequest(
@@ -106,7 +103,7 @@ public final class ProtocolRequestCreatorTest {
     @Test
     public void testCreateProtoRequest() {
         ProtocolRequestCreator requestCreator =
-                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>(), false);
+                new ProtocolRequestCreator(REQUEST_BASE_URI, new HashMap<String, String>());
 
         FederatedComputeHttpRequest request =
                 requestCreator.createProtoRequest(
diff --git a/tests/federatedcomputetests/src/com/android/federatedcompute/services/training/FederatedComputeWorkerTest.java b/tests/federatedcomputetests/src/com/android/federatedcompute/services/training/FederatedComputeWorkerTest.java
index 34179225..71fe9998 100644
--- a/tests/federatedcomputetests/src/com/android/federatedcompute/services/training/FederatedComputeWorkerTest.java
+++ b/tests/federatedcomputetests/src/com/android/federatedcompute/services/training/FederatedComputeWorkerTest.java
@@ -22,12 +22,15 @@ import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.times;
@@ -59,6 +62,7 @@ import com.android.federatedcompute.services.testutils.TrainingTestUtil;
 import com.android.federatedcompute.services.training.ResultCallbackHelper.CallbackResult;
 import com.android.federatedcompute.services.training.aidl.IIsolatedTrainingService;
 import com.android.federatedcompute.services.training.aidl.ITrainingResultCallback;
+import com.android.federatedcompute.services.training.util.ComputationResult;
 import com.android.federatedcompute.services.training.util.TrainingConditionsChecker;
 import com.android.federatedcompute.services.training.util.TrainingConditionsChecker.Condition;
 
@@ -76,15 +80,19 @@ import com.google.intelligence.fcp.client.engine.TaskRetry;
 import com.google.internal.federated.plan.ClientOnlyPlan;
 import com.google.internal.federated.plan.ClientPhase;
 import com.google.internal.federated.plan.TensorflowSpec;
+import com.google.internal.federatedcompute.v1.RejectionInfo;
+import com.google.internal.federatedcompute.v1.RetryWindow;
 import com.google.ondevicepersonalization.federatedcompute.proto.TaskAssignment;
 import com.google.protobuf.Any;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Duration;
 
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnit;
@@ -121,6 +129,15 @@ public final class FederatedComputeWorkerTest {
                     createTempFile("input", ".ckp"),
                     TrainingTestUtil.createFederatedAnalyticClientPlan(),
                     TaskAssignment.newBuilder().setTaskName(TASK_NAME).build());
+
+    public static final RejectionInfo REJECTION_INFO =
+            RejectionInfo.newBuilder()
+                    .setRetryWindow(
+                            RetryWindow.newBuilder()
+                                    .setDelayMin(Duration.newBuilder().setSeconds(3600).build())
+                                    .build())
+                    .build();
+    private static final CheckinResult REJECTION_CHECKIN_RESULT = new CheckinResult(REJECTION_INFO);
     private static final FLRunnerResult FL_RUNNER_FAILURE_RESULT =
             FLRunnerResult.newBuilder().setContributionResult(ContributionResult.FAIL).build();
 
@@ -172,9 +189,9 @@ public final class FederatedComputeWorkerTest {
     @Mock FederatedComputeJobManager mMockJobManager;
     private Context mContext;
     private FederatedComputeWorker mSpyWorker;
-    @Mock private HttpFederatedProtocol mMockHttpFederatedProtocol;
+    private HttpFederatedProtocol mSpyHttpFederatedProtocol;
     @Mock private ComputationRunner mMockComputationRunner;
-    @Mock private ResultCallbackHelper mMockResultCallbackHelper;
+    private ResultCallbackHelper mSpyResultCallbackHelper;
 
     private static byte[] createTrainingConstraints(
             boolean requiresSchedulerIdle,
@@ -201,6 +218,10 @@ public final class FederatedComputeWorkerTest {
     @Before
     public void doBeforeEachTest() throws Exception {
         mContext = ApplicationProvider.getApplicationContext();
+        mSpyHttpFederatedProtocol =
+                Mockito.spy(
+                        HttpFederatedProtocol.create(SERVER_ADDRESS, "1.0.0.1", POPULATION_NAME));
+        mSpyResultCallbackHelper = Mockito.spy(new ResultCallbackHelper(mContext));
         mSpyWorker =
                 Mockito.spy(
                         new FederatedComputeWorker(
@@ -208,14 +229,15 @@ public final class FederatedComputeWorkerTest {
                                 mMockJobManager,
                                 mTrainingConditionsChecker,
                                 mMockComputationRunner,
-                                mMockResultCallbackHelper,
+                                mSpyResultCallbackHelper,
                                 new TestInjector()));
         when(mTrainingConditionsChecker.checkAllConditionsForFlTraining(any()))
                 .thenReturn(EnumSet.noneOf(Condition.class));
-        when(mMockResultCallbackHelper.callHandleResult(eq(TASK_NAME), any(), any()))
-                .thenReturn(Futures.immediateFuture(CallbackResult.SUCCESS));
+        doReturn(Futures.immediateFuture(CallbackResult.SUCCESS))
+                .when(mSpyResultCallbackHelper)
+                .callHandleResult(eq(TASK_NAME), any(), any());
         when(mMockJobManager.onTrainingStarted(anyInt())).thenReturn(FEDERATED_TRAINING_TASK_1);
-        doReturn(mMockHttpFederatedProtocol)
+        doReturn(mSpyHttpFederatedProtocol)
                 .when(mSpyWorker)
                 .getHttpFederatedProtocol(anyString(), anyString());
         when(mMockComputationRunner.runTaskWithNativeRunner(
@@ -263,10 +285,10 @@ public final class FederatedComputeWorkerTest {
                                 new ExecutionException(
                                         "issue checkin failed",
                                         new IllegalStateException("http 404"))))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .issueCheckin();
         doReturn(FluentFuture.from(immediateFuture(null)))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .reportResult(any());
 
         assertThrows(ExecutionException.class, () -> mSpyWorker.startTrainingRun(JOB_ID).get());
@@ -278,11 +300,58 @@ public final class FederatedComputeWorkerTest {
     }
 
     @Test
+    public void testCheckinWithRejection() throws Exception {
+        setUpExampleStoreService();
+        doReturn(
+                immediateFuture(REJECTION_CHECKIN_RESULT))
+                .when(mSpyHttpFederatedProtocol)
+                .issueCheckin();
+
+        FLRunnerResult result = mSpyWorker.startTrainingRun(JOB_ID).get();
+
+        assertNull(result);
+        verify(mMockJobManager)
+                .onTrainingCompleted(
+                        anyInt(), anyString(), any(), any(), eq(ContributionResult.FAIL));
+        mSpyWorker.finish(null, ContributionResult.FAIL, false);
+    }
+
+    @Test
+    public void testReportResultWithRejection() throws Exception {
+        setUpExampleStoreService();
+        doReturn(immediateFuture(FA_CHECKIN_RESULT))
+                .when(mSpyHttpFederatedProtocol)
+                .issueCheckin();
+        doReturn(FluentFuture.from(immediateFuture(REJECTION_INFO)))
+                .when(mSpyHttpFederatedProtocol)
+                .reportResult(any());
+        doCallRealMethod().when(mSpyResultCallbackHelper).callHandleResult(any(), any(), any());
+        ArgumentCaptor<ComputationResult> computationResultCaptor =
+                ArgumentCaptor.forClass(ComputationResult.class);
+
+
+        FLRunnerResult result = mSpyWorker.startTrainingRun(JOB_ID).get();
+
+        assertNull(result);
+        verify(mMockJobManager)
+                .onTrainingCompleted(
+                        anyInt(), anyString(), any(), any(), eq(ContributionResult.FAIL));
+        verify(mSpyResultCallbackHelper)
+                .callHandleResult(any(), any(), computationResultCaptor.capture());
+        ComputationResult computationResult = computationResultCaptor.getValue();
+        assertNotNull(computationResult.getFlRunnerResult());
+        assertEquals(
+                ContributionResult.FAIL,
+                computationResult.getFlRunnerResult().getContributionResult());
+        mSpyWorker.finish(null, ContributionResult.FAIL, false);
+    }
+
+    @Test
     public void testReportResultFails_throwsException() throws Exception {
         setUpExampleStoreService();
 
         doReturn(immediateFuture(FA_CHECKIN_RESULT))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .issueCheckin();
         doReturn(
                         FluentFuture.from(
@@ -290,7 +359,7 @@ public final class FederatedComputeWorkerTest {
                                         new ExecutionException(
                                                 "report result failed",
                                                 new IllegalStateException("http 404")))))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .reportResult(any());
 
         assertThrows(ExecutionException.class, () -> mSpyWorker.startTrainingRun(JOB_ID).get());
@@ -305,10 +374,10 @@ public final class FederatedComputeWorkerTest {
     @Test
     public void testBindToExampleStoreFails_throwsException() throws Exception {
         setUpHttpFederatedProtocol(FL_CHECKIN_RESULT);
-
         // Mock failure bind to ExampleStoreService.
         doReturn(null).when(mSpyWorker).getExampleStoreService(anyString());
-        doNothing().when(mSpyWorker).unbindFromExampleStoreService();
+        ArgumentCaptor<ComputationResult> computationResultCaptor =
+                ArgumentCaptor.forClass(ComputationResult.class);
 
         assertThrows(ExecutionException.class, () -> mSpyWorker.startTrainingRun(JOB_ID).get());
 
@@ -317,6 +386,13 @@ public final class FederatedComputeWorkerTest {
                 .onTrainingCompleted(
                         anyInt(), anyString(), any(), any(), eq(ContributionResult.FAIL));
         verify(mSpyWorker, times(0)).unbindFromExampleStoreService();
+        verify(mSpyHttpFederatedProtocol, times(1))
+                .reportResult(computationResultCaptor.capture());
+        ComputationResult computationResult = computationResultCaptor.getValue();
+        assertNotNull(computationResult.getFlRunnerResult());
+        assertEquals(
+                ContributionResult.FAIL,
+                computationResult.getFlRunnerResult().getContributionResult());
     }
 
     @Test
@@ -348,14 +424,52 @@ public final class FederatedComputeWorkerTest {
     }
 
     @Test
+    public void testRunFAComputationThrows() throws Exception {
+        setUpExampleStoreService();
+        setUpHttpFederatedProtocol(FA_CHECKIN_RESULT);
+        doReturn(FluentFuture.from(immediateFuture(null)))
+                .when(mSpyHttpFederatedProtocol)
+                .reportResult(any());
+        when(mMockComputationRunner.runTaskWithNativeRunner(
+                        anyString(),
+                        anyString(),
+                        anyString(),
+                        anyString(),
+                        any(),
+                        any(),
+                        any(),
+                        any(),
+                        any()))
+                .thenThrow(new RuntimeException("Test failures!"));
+        ArgumentCaptor<ComputationResult> computationResultCaptor =
+                ArgumentCaptor.forClass(ComputationResult.class);
+
+        assertThrows(ExecutionException.class, () -> mSpyWorker.startTrainingRun(JOB_ID).get());
+
+        mSpyWorker.finish(null, ContributionResult.FAIL, false);
+        verify(mMockJobManager)
+                .onTrainingCompleted(
+                        anyInt(), anyString(), any(), any(), eq(ContributionResult.FAIL));
+        verify(mSpyWorker).unbindFromExampleStoreService();
+        verify(mSpyHttpFederatedProtocol, times(1))
+                .reportResult(computationResultCaptor.capture());
+        ComputationResult computationResult = computationResultCaptor.getValue();
+        assertNotNull(computationResult.getFlRunnerResult());
+        assertEquals(
+                ContributionResult.FAIL,
+                computationResult.getFlRunnerResult().getContributionResult());
+    }
+
+    @Test
     public void testPublishToResultHandlingServiceFails_returnsSuccess() throws Exception {
         setUpExampleStoreService();
         setUpHttpFederatedProtocol(FA_CHECKIN_RESULT);
 
         // Mock publish to ResultHandlingService fails which is best effort and should not affect
         // final result.
-        when(mMockResultCallbackHelper.callHandleResult(eq(TASK_NAME), any(), any()))
-                .thenReturn(Futures.immediateFuture(CallbackResult.FAIL));
+        doReturn(Futures.immediateFuture(CallbackResult.FAIL))
+                .when(mSpyResultCallbackHelper)
+                .callHandleResult(eq(TASK_NAME), any(), any());
 
         FLRunnerResult result = mSpyWorker.startTrainingRun(JOB_ID).get();
         assertThat(result.getContributionResult()).isEqualTo(ContributionResult.SUCCESS);
@@ -375,12 +489,12 @@ public final class FederatedComputeWorkerTest {
 
         // Mock publish to ResultHandlingService throws exception which is best effort and should
         // not affect final result.
-        when(mMockResultCallbackHelper.callHandleResult(eq(TASK_NAME), any(), any()))
-                .thenReturn(
-                        immediateFailedFuture(
-                                new ExecutionException(
-                                        "ResultHandlingService fail",
-                                        new IllegalStateException("can't bind to service"))));
+        doReturn(immediateFailedFuture(
+                new ExecutionException(
+                        "ResultHandlingService fail",
+                        new IllegalStateException("can't bind to service"))))
+                .when(mSpyResultCallbackHelper)
+                .callHandleResult(eq(TASK_NAME), any(), any());
 
         FLRunnerResult result = mSpyWorker.startTrainingRun(JOB_ID).get();
         assertThat(result.getContributionResult()).isEqualTo(ContributionResult.SUCCESS);
@@ -390,7 +504,7 @@ public final class FederatedComputeWorkerTest {
                 .onTrainingCompleted(
                         anyInt(), anyString(), any(), any(), eq(ContributionResult.SUCCESS));
         verify(mSpyWorker).unbindFromExampleStoreService();
-        verify(mMockResultCallbackHelper).callHandleResult(eq(TASK_NAME), any(), any());
+        verify(mSpyResultCallbackHelper).callHandleResult(eq(TASK_NAME), any(), any());
     }
 
     @Test
@@ -408,7 +522,7 @@ public final class FederatedComputeWorkerTest {
     @Test
     public void testBindToIsolatedTrainingServiceFail_returnsFail() throws Exception {
         doReturn(immediateFuture(FL_CHECKIN_RESULT))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .issueCheckin();
         setUpExampleStoreService();
 
@@ -489,9 +603,9 @@ public final class FederatedComputeWorkerTest {
     }
 
     private void setUpHttpFederatedProtocol(CheckinResult checkinResult) {
-        doReturn(immediateFuture(checkinResult)).when(mMockHttpFederatedProtocol).issueCheckin();
+        doReturn(immediateFuture(checkinResult)).when(mSpyHttpFederatedProtocol).issueCheckin();
         doReturn(FluentFuture.from(immediateFuture(null)))
-                .when(mMockHttpFederatedProtocol)
+                .when(mSpyHttpFederatedProtocol)
                 .reportResult(any());
     }
 
diff --git a/tests/frameworktests/src/android/adservices/ondevicepersonalization/IsolatedServiceTest.java b/tests/frameworktests/src/android/adservices/ondevicepersonalization/IsolatedServiceTest.java
index a7c11057..caef0bdb 100644
--- a/tests/frameworktests/src/android/adservices/ondevicepersonalization/IsolatedServiceTest.java
+++ b/tests/frameworktests/src/android/adservices/ondevicepersonalization/IsolatedServiceTest.java
@@ -83,8 +83,8 @@ public class IsolatedServiceTest {
     public void testOnExecute() throws Exception {
         PersistableBundle appParams = new PersistableBundle();
         appParams.putString("x", "y");
-        ExecuteInput input =
-                new ExecuteInput.Builder()
+        ExecuteInputParcel input =
+                new ExecuteInputParcel.Builder()
                         .setAppPackageName("com.testapp")
                         .setAppParams(appParams)
                         .build();
@@ -97,8 +97,8 @@ public class IsolatedServiceTest {
         mBinder.onRequest(Constants.OP_EXECUTE, params, new TestServiceCallback());
         mLatch.await();
         assertTrue(mSelectContentCalled);
-        ExecuteOutput result =
-                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, ExecuteOutput.class);
+        ExecuteOutputParcel result =
+                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, ExecuteOutputParcel.class);
         assertEquals(5, result.getRequestLogRecord().getRows().get(0).getAsInteger("a").intValue());
         assertEquals("123", result.getRenderingConfigs().get(0).getKeys().get(0));
     }
@@ -107,8 +107,8 @@ public class IsolatedServiceTest {
     public void testOnExecutePropagatesError() throws Exception {
         PersistableBundle appParams = new PersistableBundle();
         appParams.putInt("error", 1); // Trigger an error in the service.
-        ExecuteInput input =
-                new ExecuteInput.Builder()
+        ExecuteInputParcel input =
+                new ExecuteInputParcel.Builder()
                         .setAppPackageName("com.testapp")
                         .setAppParams(appParams)
                         .build();
@@ -126,7 +126,7 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnExecuteWithoutAppParams() throws Exception {
-        ExecuteInput input = new ExecuteInput.Builder().setAppPackageName("com.testapp").build();
+        ExecuteInputParcel input = new ExecuteInputParcel.Builder().setAppPackageName("com.testapp").build();
         Bundle params = new Bundle();
         params.putParcelable(Constants.EXTRA_INPUT, input);
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
@@ -163,7 +163,7 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnExecuteThrowsIfDataAccessServiceMissing() throws Exception {
-        ExecuteInput input = new ExecuteInput.Builder().setAppPackageName("com.testapp").build();
+        ExecuteInputParcel input = new ExecuteInputParcel.Builder().setAppPackageName("com.testapp").build();
         Bundle params = new Bundle();
         params.putBinder(
                 Constants.EXTRA_FEDERATED_COMPUTE_SERVICE_BINDER,
@@ -178,7 +178,7 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnExecuteThrowsIfFederatedComputeServiceMissing() throws Exception {
-        ExecuteInput input = new ExecuteInput.Builder().setAppPackageName("com.testapp").build();
+        ExecuteInputParcel input = new ExecuteInputParcel.Builder().setAppPackageName("com.testapp").build();
         Bundle params = new Bundle();
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
         params.putParcelable(Constants.EXTRA_INPUT, input);
@@ -191,7 +191,7 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnExecuteThrowsIfCallbackMissing() throws Exception {
-        ExecuteInput input = new ExecuteInput.Builder().setAppPackageName("com.testapp").build();
+        ExecuteInputParcel input = new ExecuteInputParcel.Builder().setAppPackageName("com.testapp").build();
         Bundle params = new Bundle();
         params.putParcelable(Constants.EXTRA_INPUT, input);
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
@@ -218,9 +218,9 @@ public class IsolatedServiceTest {
         mBinder.onRequest(Constants.OP_DOWNLOAD, params, new TestServiceCallback());
         mLatch.await();
         assertTrue(mOnDownloadCalled);
-        DownloadCompletedOutput result =
+        DownloadCompletedOutputParcel result =
                 mCallbackResult.getParcelable(
-                        Constants.EXTRA_RESULT, DownloadCompletedOutput.class);
+                        Constants.EXTRA_RESULT, DownloadCompletedOutputParcel.class);
         assertEquals("12", result.getRetainedKeys().get(0));
     }
 
@@ -297,8 +297,8 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnRender() throws Exception {
-        RenderInput input =
-                new RenderInput.Builder()
+        RenderInputParcel input =
+                new RenderInputParcel.Builder()
                         .setRenderingConfig(
                                 new RenderingConfig.Builder().addKey("a").addKey("b").build())
                         .build();
@@ -308,15 +308,15 @@ public class IsolatedServiceTest {
         mBinder.onRequest(Constants.OP_RENDER, params, new TestServiceCallback());
         mLatch.await();
         assertTrue(mOnRenderCalled);
-        RenderOutput result =
-                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, RenderOutput.class);
+        RenderOutputParcel result =
+                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, RenderOutputParcel.class);
         assertEquals("htmlstring", result.getContent());
     }
 
     @Test
     public void testOnRenderPropagatesError() throws Exception {
-        RenderInput input =
-                new RenderInput.Builder()
+        RenderInputParcel input =
+                new RenderInputParcel.Builder()
                         .setRenderingConfig(
                                 new RenderingConfig.Builder()
                                         .addKey("z") // Trigger error in service.
@@ -353,8 +353,8 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnRenderThrowsIfDataAccessServiceMissing() throws Exception {
-        RenderInput input =
-                new RenderInput.Builder()
+        RenderInputParcel input =
+                new RenderInputParcel.Builder()
                         .setRenderingConfig(
                                 new RenderingConfig.Builder().addKey("a").addKey("b").build())
                         .build();
@@ -369,8 +369,8 @@ public class IsolatedServiceTest {
 
     @Test
     public void testOnRenderThrowsIfCallbackMissing() throws Exception {
-        RenderInput input =
-                new RenderInput.Builder()
+        RenderInputParcel input =
+                new RenderInputParcel.Builder()
                         .setRenderingConfig(
                                 new RenderingConfig.Builder().addKey("a").addKey("b").build())
                         .build();
@@ -385,29 +385,30 @@ public class IsolatedServiceTest {
     }
 
     @Test
-    public void testOnWebViewEvent() throws Exception {
+    public void testOnEvent() throws Exception {
         Bundle params = new Bundle();
         params.putParcelable(
                 Constants.EXTRA_INPUT,
-                new EventInput.Builder().setParameters(PersistableBundle.EMPTY).build());
+                new EventInputParcel.Builder().setParameters(PersistableBundle.EMPTY).build());
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
         mBinder.onRequest(Constants.OP_WEB_VIEW_EVENT, params, new TestServiceCallback());
         mLatch.await();
         assertTrue(mOnEventCalled);
-        EventOutput result =
-                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, EventOutput.class);
+        EventOutputParcel result =
+                mCallbackResult.getParcelable(Constants.EXTRA_RESULT, EventOutputParcel.class);
         assertEquals(1, result.getEventLogRecord().getType());
         assertEquals(2, result.getEventLogRecord().getRowIndex());
     }
 
     @Test
-    public void testOnWebViewEventPropagatesError() throws Exception {
+    public void testOnEventPropagatesError() throws Exception {
         PersistableBundle eventParams = new PersistableBundle();
         // Input value 9999 will trigger an error in the mock service.
         eventParams.putInt(EVENT_TYPE_KEY, 9999);
         Bundle params = new Bundle();
         params.putParcelable(
-                Constants.EXTRA_INPUT, new EventInput.Builder().setParameters(eventParams).build());
+                Constants.EXTRA_INPUT,
+                new EventInputParcel.Builder().setParameters(eventParams).build());
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
         mBinder.onRequest(Constants.OP_WEB_VIEW_EVENT, params, new TestServiceCallback());
         mLatch.await();
@@ -441,7 +442,7 @@ public class IsolatedServiceTest {
         Bundle params = new Bundle();
         params.putParcelable(
                 Constants.EXTRA_INPUT,
-                new EventInput.Builder().setParameters(PersistableBundle.EMPTY).build());
+                new EventInputParcel.Builder().setParameters(PersistableBundle.EMPTY).build());
         assertThrows(
                 NullPointerException.class,
                 () -> {
@@ -455,7 +456,7 @@ public class IsolatedServiceTest {
         Bundle params = new Bundle();
         params.putParcelable(
                 Constants.EXTRA_INPUT,
-                new EventInput.Builder().setParameters(PersistableBundle.EMPTY).build());
+                new EventInputParcel.Builder().setParameters(PersistableBundle.EMPTY).build());
         params.putBinder(Constants.EXTRA_DATA_ACCESS_SERVICE_BINDER, new TestDataAccessService());
         assertThrows(
                 NullPointerException.class,
@@ -465,10 +466,10 @@ public class IsolatedServiceTest {
     }
 
     @Test
-    public void testOnTrainingExample() throws Exception {
+    public void testOnTrainingExamples() throws Exception {
         JoinedLogRecord joinedLogRecord = new JoinedLogRecord.Builder().build();
-        TrainingExampleInput input =
-                new TrainingExampleInput.Builder()
+        TrainingExamplesInputParcel input =
+                new TrainingExamplesInputParcel.Builder()
                         .setPopulationName("")
                         .setTaskName("")
                         .setResumptionToken(new byte[] {0})
@@ -479,9 +480,9 @@ public class IsolatedServiceTest {
         mBinder.onRequest(Constants.OP_TRAINING_EXAMPLE, params, new TestServiceCallback());
         mLatch.await();
         assertTrue(mOnTrainingExampleCalled);
-        TrainingExampleOutputParcel result =
+        TrainingExamplesOutputParcel result =
                 mCallbackResult.getParcelable(
-                        Constants.EXTRA_RESULT, TrainingExampleOutputParcel.class);
+                        Constants.EXTRA_RESULT, TrainingExamplesOutputParcel.class);
         List<byte[]> examples = result.getTrainingExamples().getList();
         List<byte[]> tokens = result.getResumptionTokens().getList();
         assertEquals(1, examples.size());
@@ -503,8 +504,8 @@ public class IsolatedServiceTest {
     @Test
     public void testOnTrainingExampleThrowsIfDataAccessServiceMissing() throws Exception {
         JoinedLogRecord joinedLogRecord = new JoinedLogRecord.Builder().build();
-        TrainingExampleInput input =
-                new TrainingExampleInput.Builder()
+        TrainingExamplesInputParcel input =
+                new TrainingExamplesInputParcel.Builder()
                         .setPopulationName("")
                         .setTaskName("")
                         .setResumptionToken(new byte[] {0})
@@ -522,8 +523,8 @@ public class IsolatedServiceTest {
     @Test
     public void testOnTrainingExampleThrowsIfCallbackMissing() throws Exception {
         JoinedLogRecord joinedLogRecord = new JoinedLogRecord.Builder().build();
-        TrainingExampleInput input =
-                new TrainingExampleInput.Builder()
+        TrainingExamplesInputParcel input =
+                new TrainingExamplesInputParcel.Builder()
                         .setPopulationName("")
                         .setTaskName("")
                         .setResumptionToken(new byte[] {0})
@@ -608,15 +609,15 @@ public class IsolatedServiceTest {
         }
 
         @Override
-        public void onTrainingExample(
-                TrainingExampleInput input, Consumer<TrainingExampleOutput> consumer) {
+        public void onTrainingExamples(
+                TrainingExamplesInput input, Consumer<TrainingExamplesOutput> consumer) {
             mOnTrainingExampleCalled = true;
             List<byte[]> examples = new ArrayList<>();
             examples.add(new byte[] {12});
             List<byte[]> tokens = new ArrayList<>();
             tokens.add(new byte[] {13});
             consumer.accept(
-                    new TrainingExampleOutput.Builder()
+                    new TrainingExamplesOutput.Builder()
                             .setTrainingExamples(examples)
                             .setResumptionTokens(tokens)
                             .build());
diff --git a/tests/frameworktests/src/android/adservices/ondevicepersonalization/LogReaderTest.java b/tests/frameworktests/src/android/adservices/ondevicepersonalization/LogReaderTest.java
index 6397a44e..e0056d0a 100644
--- a/tests/frameworktests/src/android/adservices/ondevicepersonalization/LogReaderTest.java
+++ b/tests/frameworktests/src/android/adservices/ondevicepersonalization/LogReaderTest.java
@@ -35,6 +35,7 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -56,7 +57,8 @@ public class LogReaderTest {
 
     @Test
     public void testGetRequestsSuccess() {
-        List<RequestLogRecord> result = mLogReader.getRequests(10, 100);
+        List<RequestLogRecord> result = mLogReader.getRequests(
+                Instant.ofEpochMilli(10), Instant.ofEpochMilli(100));
         assertEquals(2, result.size());
         assertEquals(1, result.get(0).getRows().size());
         assertEquals((int) (result.get(0).getRows().get(0).getAsInteger("a")), 1);
@@ -67,25 +69,38 @@ public class LogReaderTest {
     }
 
     @Test
+    public void testGetRequestsNullTimeError() {
+        assertThrows(NullPointerException.class, () -> mLogReader.getRequests(
+                null, Instant.ofEpochMilli(100)));
+        assertThrows(NullPointerException.class, () -> mLogReader.getRequests(
+                Instant.ofEpochMilli(100), null));
+    }
+
+    @Test
     public void testGetRequestsError() {
         // Triggers an expected error in the mock service.
-        assertThrows(IllegalStateException.class, () -> mLogReader.getRequests(7, 100));
+        assertThrows(IllegalStateException.class, () -> mLogReader.getRequests(
+                Instant.ofEpochMilli(7), Instant.ofEpochMilli(100)));
     }
 
     @Test
     public void testGetRequestsNegativeTimeError() {
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(-1, 100));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(
+                Instant.ofEpochMilli(-1), Instant.ofEpochMilli(100)));
     }
 
     @Test
     public void testGetRequestsBadTimeRangeError() {
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(100, 100));
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(1000, 100));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(
+                Instant.ofEpochMilli(100), Instant.ofEpochMilli(100)));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getRequests(
+                Instant.ofEpochMilli(1000), Instant.ofEpochMilli(100)));
     }
 
     @Test
     public void testGetJoinedEventsSuccess() {
-        List<EventLogRecord> result = mLogReader.getJoinedEvents(10, 100);
+        List<EventLogRecord> result = mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(10), Instant.ofEpochMilli(100));
         assertEquals(2, result.size());
         assertEquals(result.get(0).getTimeMillis(), 30);
         assertEquals(result.get(0).getRequestLogRecord().getTimeMillis(), 20);
@@ -102,18 +117,30 @@ public class LogReaderTest {
     @Test
     public void testGetJoinedEventsError() {
         // Triggers an expected error in the mock service.
-        assertThrows(IllegalStateException.class, () -> mLogReader.getJoinedEvents(7, 100));
+        assertThrows(IllegalStateException.class, () -> mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(7), Instant.ofEpochMilli(100)));
+    }
+
+    @Test
+    public void testGetJoinedEventsNullTimeError() {
+        assertThrows(NullPointerException.class, () -> mLogReader.getJoinedEvents(
+                null, Instant.ofEpochMilli(100)));
+        assertThrows(NullPointerException.class, () -> mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(100), null));
     }
 
     @Test
     public void testGetJoinedEventsNegativeTimeError() {
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(-1, 100));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(-1), Instant.ofEpochMilli(100)));
     }
 
     @Test
     public void testGetJoinedEventsInputError() {
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(100, 100));
-        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(1000, 100));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(100), Instant.ofEpochMilli(100)));
+        assertThrows(IllegalArgumentException.class, () -> mLogReader.getJoinedEvents(
+                Instant.ofEpochMilli(1000), Instant.ofEpochMilli(100)));
     }
 
     public static class LocalDataService extends IDataAccessService.Stub {
diff --git a/tests/frameworktests/src/android/adservices/ondevicepersonalization/OnDevicePersonalizationFrameworkClassesTest.java b/tests/frameworktests/src/android/adservices/ondevicepersonalization/OnDevicePersonalizationFrameworkClassesTest.java
index eac96f6c..258ecac1 100644
--- a/tests/frameworktests/src/android/adservices/ondevicepersonalization/OnDevicePersonalizationFrameworkClassesTest.java
+++ b/tests/frameworktests/src/android/adservices/ondevicepersonalization/OnDevicePersonalizationFrameworkClassesTest.java
@@ -37,23 +37,46 @@ import java.util.ArrayList;
 @RunWith(AndroidJUnit4.class)
 public class OnDevicePersonalizationFrameworkClassesTest {
     /**
+     * Tests that the ExecuteInput object serializes correctly.
+     */
+    @Test
+    public void testExecuteInput() {
+        PersistableBundle bundle = new PersistableBundle();
+        bundle.putInt("a", 5);
+        ExecuteInputParcel data = new ExecuteInputParcel.Builder()
+                .setAppPackageName("com.example.test")
+                .setAppParams(bundle)
+                .build();
+
+        Parcel parcel = Parcel.obtain();
+        data.writeToParcel(parcel, 0);
+        parcel.setDataPosition(0);
+        ExecuteInputParcel data2 = ExecuteInputParcel.CREATOR.createFromParcel(parcel);
+        ExecuteInput result = new ExecuteInput(data2);
+
+        assertEquals("com.example.test", result.getAppPackageName());
+        assertEquals(5, result.getAppParams().getInt("a"));
+    }
+
+    /**
      * Tests that the ExecuteOutput object serializes correctly.
      */
     @Test
     public void testExecuteOutput() {
         ContentValues row = new ContentValues();
         row.put("a", 5);
-        ExecuteOutput result =
+        ExecuteOutput data =
                 new ExecuteOutput.Builder()
                     .setRequestLogRecord(new RequestLogRecord.Builder().addRow(row).build())
                     .addRenderingConfig(new RenderingConfig.Builder().addKey("abc").build())
                     .addEventLogRecord(new EventLogRecord.Builder().setType(1).build())
                     .build();
+        ExecuteOutputParcel result = new ExecuteOutputParcel(data);
 
         Parcel parcel = Parcel.obtain();
         result.writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        ExecuteOutput result2 = ExecuteOutput.CREATOR.createFromParcel(parcel);
+        ExecuteOutputParcel result2 = ExecuteOutputParcel.CREATOR.createFromParcel(parcel);
 
         assertEquals(
                 5, result2.getRequestLogRecord().getRows().get(0).getAsInteger("a").intValue());
@@ -62,16 +85,41 @@ public class OnDevicePersonalizationFrameworkClassesTest {
     }
 
     /**
+     * Tests that the RenderInput object serializes correctly.
+     */
+    @Test
+    public void testRenderInput() {
+        RenderInputParcel data = new RenderInputParcel.Builder()
+                .setWidth(10)
+                .setHeight(20)
+                .setRenderingConfigIndex(5)
+                .setRenderingConfig(new RenderingConfig.Builder().addKey("abc").build())
+                .build();
+
+        Parcel parcel = Parcel.obtain();
+        data.writeToParcel(parcel, 0);
+        parcel.setDataPosition(0);
+        RenderInputParcel data2 = RenderInputParcel.CREATOR.createFromParcel(parcel);
+        RenderInput result = new RenderInput(data2);
+
+        assertEquals(10, result.getWidth());
+        assertEquals(20, result.getHeight());
+        assertEquals(5, result.getRenderingConfigIndex());
+        assertEquals("abc", result.getRenderingConfig().getKeys().get(0));
+    }
+
+    /**
      * Tests that the RenderOutput object serializes correctly.
      */
     @Test
     public void testRenderOutput() {
-        RenderOutput result = new RenderOutput.Builder().setContent("abc").build();
+        RenderOutput data = new RenderOutput.Builder().setContent("abc").build();
+        RenderOutputParcel result = new RenderOutputParcel(data);
 
         Parcel parcel = Parcel.obtain();
         result.writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        RenderOutput result2 = RenderOutput.CREATOR.createFromParcel(parcel);
+        RenderOutputParcel result2 = RenderOutputParcel.CREATOR.createFromParcel(parcel);
 
         assertEquals("abc", result2.getContent());
     }
@@ -81,16 +129,17 @@ public class OnDevicePersonalizationFrameworkClassesTest {
      */
     @Test
     public void teetDownloadCompletedOutput() {
-        DownloadCompletedOutput result = new DownloadCompletedOutput.Builder()
+        DownloadCompletedOutput data = new DownloadCompletedOutput.Builder()
                 .addRetainedKey("abc").addRetainedKey("def").build();
+        DownloadCompletedOutputParcel result =
+                new DownloadCompletedOutputParcel(data);
 
         Parcel parcel = Parcel.obtain();
         result.writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        DownloadCompletedOutput result2 =
-                DownloadCompletedOutput.CREATOR.createFromParcel(parcel);
+        DownloadCompletedOutputParcel result2 =
+                DownloadCompletedOutputParcel.CREATOR.createFromParcel(parcel);
 
-        assertEquals(result, result2);
         assertEquals("abc", result2.getRetainedKeys().get(0));
         assertEquals("def", result2.getRetainedKeys().get(1));
     }
@@ -105,7 +154,7 @@ public class OnDevicePersonalizationFrameworkClassesTest {
         ArrayList<ContentValues> rows = new ArrayList<>();
         rows.add(new ContentValues());
         rows.get(0).put("a", 5);
-        EventInput result = new EventInput.Builder()
+        EventInputParcel data = new EventInputParcel.Builder()
                 .setParameters(params)
                 .setRequestLogRecord(
                     new RequestLogRecord.Builder()
@@ -114,13 +163,14 @@ public class OnDevicePersonalizationFrameworkClassesTest {
                 .build();
 
         Parcel parcel = Parcel.obtain();
-        result.writeToParcel(parcel, 0);
+        data.writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        EventInput result2 = EventInput.CREATOR.createFromParcel(parcel);
+        EventInputParcel data2 = EventInputParcel.CREATOR.createFromParcel(parcel);
+        EventInput result = new EventInput(data2);
 
-        assertEquals(3, result2.getParameters().getInt("x"));
+        assertEquals(3, result.getParameters().getInt("x"));
         assertEquals(
-                5, result2.getRequestLogRecord().getRows().get(0).getAsInteger("a").intValue());
+                5, result.getRequestLogRecord().getRows().get(0).getAsInteger("a").intValue());
     }
 
     /**
@@ -130,7 +180,7 @@ public class OnDevicePersonalizationFrameworkClassesTest {
     public void testEventOutput() {
         ContentValues data = new ContentValues();
         data.put("a", 3);
-        EventOutput result = new EventOutput.Builder()
+        EventOutput output = new EventOutput.Builder()
                 .setEventLogRecord(
                     new EventLogRecord.Builder()
                         .setType(5)
@@ -138,13 +188,13 @@ public class OnDevicePersonalizationFrameworkClassesTest {
                         .setData(data)
                         .build())
                 .build();
+        EventOutputParcel result = new EventOutputParcel(output);
 
         Parcel parcel = Parcel.obtain();
         result.writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        EventOutput result2 = EventOutput.CREATOR.createFromParcel(parcel);
+        EventOutputParcel result2 = EventOutputParcel.CREATOR.createFromParcel(parcel);
 
-        assertEquals(result, result2);
         assertEquals(5, result2.getEventLogRecord().getType());
         assertEquals(6, result2.getEventLogRecord().getRowIndex());
         assertEquals(3, result2.getEventLogRecord().getData().getAsInteger("a").intValue());
diff --git a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleAndForceTraining.java b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleAndForceTraining.java
index d66ca093..61466b5a 100644
--- a/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleAndForceTraining.java
+++ b/tests/perftests/scenarios/src/android/federatedcompute/test/scenario/federatedcompute/ScheduleAndForceTraining.java
@@ -37,8 +37,9 @@ public class ScheduleAndForceTraining {
 
     /** Prepare the device before entering the test class */
     @BeforeClass
-    public static void prepareDevice() {
+    public static void prepareDevice() throws IOException {
         TestHelper.initialize();
+        TestHelper.killRunningProcess();
     }
 
     @Before
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 1a4fc6a3..c8ff1f88 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
@@ -70,6 +70,34 @@ public class TestHelper {
         executeShellCommand(
                 "am broadcast -a android.intent.action.BOOT_COMPLETED -p "
                     + "com.google.android.federatedcompute");
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1000");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1006");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1003");
+        SystemClock.sleep(5000);
+        executeShellCommand(
+                "cmd jobscheduler run -f "
+                        + "com.google.android.ondevicepersonalization.services 1004");
+        SystemClock.sleep(5000);
+    }
+
+    /** Kill running processes to get performance measurement under cold start */
+    public static void killRunningProcess() throws IOException {
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services");
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services:"
+                + "com.android.ondevicepersonalization."
+                + "libraries.plugin.internal.PluginExecutorService");
+        executeShellCommand("am kill com.google.android.federatedcompute");
+        executeShellCommand("am kill com.google.android.federatedcompute:"
+                + "com.android.federatedcompute.services.training.IsolatedTrainingService");
+        SystemClock.sleep(2000);
     }
 
     /** Commands to return device to original state */
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 1bda1e99..d181f82f 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
@@ -73,6 +73,15 @@ public class DownloadHelper {
                     + "com.google.android.ondevicepersonalization.services 1006");
         SystemClock.sleep(5000);
     }
+
+    /** Kill running processes to get performance measurement under cold start */
+    public static void killRunningProcess() throws IOException {
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services");
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services:"
+                + "com.android.ondevicepersonalization."
+                + "libraries.plugin.internal.PluginExecutorService");
+        SystemClock.sleep(2000);
+    }
     public static void pressHome() {
         getUiDevice().pressHome();
     }
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 a97a9fc2..2a9f3526 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
@@ -36,6 +36,7 @@ public class DownloadVendorData {
     @BeforeClass
     public static void prepareDevice() throws IOException {
         DownloadHelper.initialize();
+        DownloadHelper.killRunningProcess();
     }
     @Before
     public void setUp() throws IOException {
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversion.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversion.java
new file mode 100644
index 00000000..32b72c7b
--- /dev/null
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversion.java
@@ -0,0 +1,65 @@
+/*
+ * 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.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)
+public class ReportConversion {
+
+    private TestAppHelper mTestAppHelper = new TestAppHelper();
+
+    /** Prepare the device before entering the test class */
+    @BeforeClass
+    public static void prepareDevice() throws IOException {
+        TestAppHelper.initialize();
+        TestAppHelper.killRunningProcess();
+    }
+
+    @Before
+    public void setup() throws IOException {
+        mTestAppHelper.openApp();
+    }
+
+    @Test
+    public void testReportConversion() throws IOException {
+        mTestAppHelper.clickGetAd();
+        mTestAppHelper.verifyRenderedView();
+        mTestAppHelper.clickAd("Google!");
+        mTestAppHelper.openApp();
+        mTestAppHelper.inputSourceAdId("ad1");
+        mTestAppHelper.resetLogBuffer();
+        mTestAppHelper.clickReportConversion();
+        mTestAppHelper.verifyConversionReported();
+    }
+
+    /** Return device to original state after test exeuction */
+    @AfterClass
+    public static void tearDown() throws IOException {
+        TestAppHelper.goToHomeScreen();
+        TestAppHelper.wrapUp();
+    }
+}
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAd.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAd.java
index 71f22bca..cba70b28 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAd.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAd.java
@@ -36,6 +36,7 @@ public class RequestAd {
     @BeforeClass
     public static void prepareDevice() throws IOException {
         TestAppHelper.initialize();
+        TestAppHelper.killRunningProcess();
     }
 
     @Before
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdAndClickAd.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdAndClickAd.java
index 9cd3d359..6b2a878c 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdAndClickAd.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdAndClickAd.java
@@ -36,6 +36,7 @@ public class RequestAdAndClickAd {
     @BeforeClass
     public static void prepareDevice() throws IOException {
         TestAppHelper.initialize();
+        TestAppHelper.killRunningProcess();
     }
 
     @Before
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdWithTestAppRotations.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdWithTestAppRotations.java
index 5e899f4d..5e0b5e0f 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdWithTestAppRotations.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/RequestAdWithTestAppRotations.java
@@ -37,6 +37,7 @@ public class RequestAdWithTestAppRotations {
     @BeforeClass
     public static void prepareDevice() throws IOException {
         TestAppHelper.initialize();
+        TestAppHelper.killRunningProcess();
     }
 
     @Before
@@ -49,25 +50,32 @@ public class RequestAdWithTestAppRotations {
         mTestAppHelper.clickGetAd();
         mTestAppHelper.verifyRenderedView();
 
-        // Rotate to landscape layout
+        // Rotate to landscape layout and get Ad
         mTestAppHelper.setOrientationLandscape();
         mTestAppHelper.clickGetAd();
         mTestAppHelper.verifyRenderedView();
 
-        // Rotate to portrait layout
+        // Rotate to portrait layout and get Ad
         mTestAppHelper.setOrientationPortrait();
         mTestAppHelper.clickGetAd();
         mTestAppHelper.verifyRenderedView();
 
-        // Rotate to landscape layout
+        // Rotate to landscape layout and do nothing
         mTestAppHelper.setOrientationLandscape();
-        mTestAppHelper.clickGetAd();
-        mTestAppHelper.verifyRenderedView();
 
-        // Rotate to portrait layout
+        // Rotate to portrait layout and do nothing
         mTestAppHelper.setOrientationPortrait();
+
+        // Rotate to landscape layout and get Ad
+        mTestAppHelper.setOrientationLandscape();
         mTestAppHelper.clickGetAd();
         mTestAppHelper.verifyRenderedView();
+
+        // Rotate to portrait layout and do nothing
+        mTestAppHelper.setOrientationPortrait();
+
+        // Rotate to landscape layout and do nothing
+        mTestAppHelper.setOrientationLandscape();
     }
 
     /** Return device to original state after test exeuction */
diff --git a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/TestAppHelper.java b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/TestAppHelper.java
index 6c2c720e..0b62b3e1 100644
--- a/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/TestAppHelper.java
+++ b/tests/perftests/scenarios/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/TestAppHelper.java
@@ -21,6 +21,7 @@ import static org.junit.Assert.assertNotNull;
 
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.util.Log;
 
 import androidx.test.uiautomator.By;
 import androidx.test.uiautomator.UiDevice;
@@ -36,13 +37,19 @@ import java.io.IOException;
 
 /** Helper class for interacting with OdpClient test app in perf tests. */
 public class TestAppHelper {
+    private static final String TAG = TestAppHelper.class.getSimpleName();
     private static final UiDevice sUiDevice = UiDevice.getInstance(getInstrumentation());
     private static UiScrollable sUiScrollable;
     private static final long UI_FIND_RESOURCE_TIMEOUT = 5000;
-    private static final long UI_ROTATE_IDLE_TIMEOUT = 5000;
+    private static final long UI_ROTATE_IDLE_TIMEOUT = 2500;
     private static final String ODP_CLIENT_TEST_APP_PACKAGE_NAME = "com.example.odpclient";
     private static final String GET_AD_BUTTON_RESOURCE_ID = "get_ad_button";
     private static final String RENDERED_VIEW_RESOURCE_ID = "rendered_view";
+    private static final String REPORT_CONVERSION_TEXT_BOX_RESOURCE_ID =
+            "report_conversion_text_box";
+    private static final String REPORT_CONVERSION_BUTTON_RESOURCE_ID = "report_conversion_button";
+    private static final String REPORT_CONVERSION_SUCCESS_LOG = "execute() success";
+    private static final long REPORT_CONVERSION_TIMEOUT = 10_000;
 
     /** Commands to prepare the device and odp module before testing. */
     public static void initialize() throws IOException {
@@ -81,6 +88,15 @@ public class TestAppHelper {
         SystemClock.sleep(5000);
     }
 
+    /** Kill running processes to get performance measurement under cold start */
+    public static void killRunningProcess() throws IOException {
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services");
+        executeShellCommand("am kill com.google.android.ondevicepersonalization.services:"
+                + "com.android.ondevicepersonalization."
+                + "libraries.plugin.internal.PluginExecutorService");
+        SystemClock.sleep(2000);
+    }
+
     /** Commands to return device to original state */
     public static void wrapUp() throws IOException {
         executeShellCommand(
@@ -108,6 +124,7 @@ public class TestAppHelper {
 
     /** Rotate screen to landscape orientation */
     public void setOrientationLandscape() throws RemoteException {
+        Log.d(TAG, "Rotating screen to landscape orientation");
         sUiDevice.unfreezeRotation();
         sUiDevice.setOrientationLandscape();
         SystemClock.sleep(UI_ROTATE_IDLE_TIMEOUT);
@@ -115,6 +132,7 @@ public class TestAppHelper {
 
     /** Rotate screen to portrait orientation */
     public void setOrientationPortrait() throws RemoteException {
+        Log.d(TAG, "Rotating screen to portrait orientation");
         sUiDevice.unfreezeRotation();
         sUiDevice.setOrientationPortrait();
         SystemClock.sleep(UI_ROTATE_IDLE_TIMEOUT);
@@ -124,6 +142,7 @@ public class TestAppHelper {
     public void clickGetAd() {
         UiObject2 getAdButton = getGetAdButton();
         assertNotNull("Get Ad button not found", getAdButton);
+        Log.d(TAG, "Clicking Get Ad button");
         getAdButton.click();
     }
 
@@ -135,6 +154,8 @@ public class TestAppHelper {
         SystemClock.sleep(UI_FIND_RESOURCE_TIMEOUT);
         if (renderedView.getChildCount() == 0) {
             Assert.fail("Failed to render child surface view");
+        } else {
+            Log.d(TAG, "Verified child view is rendered");
         }
     }
 
@@ -151,6 +172,55 @@ public class TestAppHelper {
         }
     }
 
+    /** Put sourceAdId name down to report conversion */
+    public void inputSourceAdId(final String sourceAdId) {
+        UiObject2 reportConversionTextBox = getReportConversionTextBox();
+        assertNotNull("Report Conversion text box not found", reportConversionTextBox);
+        reportConversionTextBox.setText(sourceAdId);
+    }
+
+    /** Click Report Conversion button */
+    public void clickReportConversion() {
+        UiObject2 reportConversionButton = getReportConversionButton();
+        assertNotNull("Report Conversion button not found", reportConversionButton);
+        reportConversionButton.click();
+        SystemClock.sleep(3000);
+    }
+
+    /** Verify conversion is reported */
+    public void verifyConversionReported() throws IOException {
+        boolean foundReportConversionSuccessLog = findLog(
+                REPORT_CONVERSION_SUCCESS_LOG,
+                REPORT_CONVERSION_TIMEOUT,
+                2000);
+
+        if (!foundReportConversionSuccessLog) {
+            Assert.fail(String.format(
+                    "Failed to find report conversion success log within test window %d ms",
+                    REPORT_CONVERSION_TIMEOUT));
+        }
+    }
+
+    /** Clean log buffer and set a high buffer limit to capture logs for test verification */
+    public void resetLogBuffer() {
+        executeShellCommand("logcat -c"); // Cleans the log buffer
+        executeShellCommand("logcat -G 32M"); // Set log buffer to 32MB
+    }
+
+    /** 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 (sUiDevice.executeShellCommand("logcat -d").contains(targetLog)) {
+                return true;
+            }
+            SystemClock.sleep(queryIntervalMillis);
+        }
+        return false;
+    }
+
     private UiObject2 getGetAdButton() {
         return sUiDevice.wait(
             Until.findObject(By.res(ODP_CLIENT_TEST_APP_PACKAGE_NAME, GET_AD_BUTTON_RESOURCE_ID)),
@@ -185,6 +255,20 @@ public class TestAppHelper {
             UI_FIND_RESOURCE_TIMEOUT);
     }
 
+    private UiObject2 getReportConversionTextBox() {
+        return sUiDevice.wait(
+                Until.findObject(By.res(
+                        ODP_CLIENT_TEST_APP_PACKAGE_NAME, REPORT_CONVERSION_TEXT_BOX_RESOURCE_ID)),
+                UI_FIND_RESOURCE_TIMEOUT);
+    }
+
+    private UiObject2 getReportConversionButton() {
+        return sUiDevice.wait(
+                Until.findObject(By.res(
+                        ODP_CLIENT_TEST_APP_PACKAGE_NAME, REPORT_CONVERSION_BUTTON_RESOURCE_ID)),
+                UI_FIND_RESOURCE_TIMEOUT);
+    }
+
     /** Get a UiScrollable instance configured for vertical scrolling */
     private static UiScrollable getUiScrollable() {
         if (sUiScrollable == null) {
diff --git a/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversionMicrobenchmark.java b/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversionMicrobenchmark.java
new file mode 100644
index 00000000..024436ea
--- /dev/null
+++ b/tests/perftests/scenarios/tests/src/android/ondevicepersonalization/test/scenario/ondevicepersonalization/ReportConversionMicrobenchmark.java
@@ -0,0 +1,36 @@
+/*
+ * 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.KillAppsRule;
+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 ReportConversionMicrobenchmark extends ReportConversion {
+
+    @Rule
+    public RuleChain rules = RuleChain.outerRule(new DropCachesRule())
+            .around(new KillAppsRule("com.example.odpclient"))
+            .around(new KillAppsRule("com.android.chrome"))
+            .around(new PressHomeRule());
+}
diff --git a/tests/servicetests/src/com/android/ondevicepersonalization/services/display/DisplayHelperTest.java b/tests/servicetests/src/com/android/ondevicepersonalization/services/display/DisplayHelperTest.java
index b092c7a1..cd09ce69 100644
--- a/tests/servicetests/src/com/android/ondevicepersonalization/services/display/DisplayHelperTest.java
+++ b/tests/servicetests/src/com/android/ondevicepersonalization/services/display/DisplayHelperTest.java
@@ -23,6 +23,7 @@ import static org.junit.Assert.assertNotNull;
 
 import android.Manifest;
 import android.adservices.ondevicepersonalization.RenderOutput;
+import android.adservices.ondevicepersonalization.RenderOutputParcel;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.content.Context;
 import android.hardware.display.DisplayManager;
@@ -72,7 +73,8 @@ public class DisplayHelperTest {
         DisplayHelper displayHelper = new DisplayHelper(mContext);
         RenderOutput renderContentResult = new RenderOutput.Builder()
                 .setContent("html").build();
-        assertEquals("html", displayHelper.generateHtml(renderContentResult,
+        RenderOutputParcel resultParcel = new RenderOutputParcel(renderContentResult);
+        assertEquals("html", displayHelper.generateHtml(resultParcel,
                 mContext.getPackageName()));
     }
 
@@ -95,8 +97,9 @@ public class DisplayHelperTest {
                 .setTemplateId("templateId")
                 .setTemplateParams(bundle)
                 .build();
+        RenderOutputParcel resultParcel = new RenderOutputParcel(renderContentResult);
         String expected = "Hello odp! I am 100.";
-        assertEquals(expected, displayHelper.generateHtml(renderContentResult,
+        assertEquals(expected, displayHelper.generateHtml(resultParcel,
                 mContext.getPackageName()));
     }
 
diff --git a/tests/servicetests/src/com/android/ondevicepersonalization/services/request/RenderFlowTest.java b/tests/servicetests/src/com/android/ondevicepersonalization/services/request/RenderFlowTest.java
index e3b04486..61e2d373 100644
--- a/tests/servicetests/src/com/android/ondevicepersonalization/services/request/RenderFlowTest.java
+++ b/tests/servicetests/src/com/android/ondevicepersonalization/services/request/RenderFlowTest.java
@@ -20,7 +20,7 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import android.adservices.ondevicepersonalization.Constants;
-import android.adservices.ondevicepersonalization.RenderOutput;
+import android.adservices.ondevicepersonalization.RenderOutputParcel;
 import android.adservices.ondevicepersonalization.RenderingConfig;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
 import android.adservices.ondevicepersonalization.aidl.IRequestSurfacePackageCallback;
@@ -156,7 +156,8 @@ public class RenderFlowTest {
             super(mContext);
         }
 
-        @Override public String generateHtml(RenderOutput renderContentResult, String packageName) {
+        @Override public String generateHtml(
+                RenderOutputParcel renderContentResult, String packageName) {
             mRenderedContent = renderContentResult.getContent();
             mGenerateHtmlCalled = true;
             return mRenderedContent;
diff --git a/tests/servicetests/src/com/test/TestPersonalizationHandler.java b/tests/servicetests/src/com/test/TestPersonalizationHandler.java
index bd629c76..9ea28b87 100644
--- a/tests/servicetests/src/com/test/TestPersonalizationHandler.java
+++ b/tests/servicetests/src/com/test/TestPersonalizationHandler.java
@@ -29,8 +29,8 @@ import android.adservices.ondevicepersonalization.RenderInput;
 import android.adservices.ondevicepersonalization.RenderOutput;
 import android.adservices.ondevicepersonalization.RenderingConfig;
 import android.adservices.ondevicepersonalization.RequestLogRecord;
-import android.adservices.ondevicepersonalization.TrainingExampleInput;
-import android.adservices.ondevicepersonalization.TrainingExampleOutput;
+import android.adservices.ondevicepersonalization.TrainingExamplesInput;
+import android.adservices.ondevicepersonalization.TrainingExamplesOutput;
 import android.annotation.NonNull;
 import android.content.ContentValues;
 import android.util.Log;
@@ -140,10 +140,10 @@ public class TestPersonalizationHandler implements IsolatedWorker {
     }
 
     @Override
-    public void onTrainingExample(
-            @NonNull TrainingExampleInput input,
-            @NonNull Consumer<TrainingExampleOutput> consumer) {
-        Log.d(TAG, "onTrainingExample() started.");
+    public void onTrainingExamples(
+            @NonNull TrainingExamplesInput input,
+            @NonNull Consumer<TrainingExamplesOutput> consumer) {
+        Log.d(TAG, "onTrainingExamples() started.");
         Log.d(TAG, "Population name: " + input.getPopulationName());
         Log.d(TAG, "Task name: " + input.getTaskName());
 
@@ -154,8 +154,8 @@ public class TestPersonalizationHandler implements IsolatedWorker {
         tokens.add("token1".getBytes());
         tokens.add("token2".getBytes());
 
-        TrainingExampleOutput output =
-                new TrainingExampleOutput.Builder()
+        TrainingExamplesOutput output =
+                new TrainingExamplesOutput.Builder()
                         .setTrainingExamples(examples)
                         .setResumptionTokens(tokens)
                         .build();