Merge "Move BackupAgent connection to a helper class" into main
diff --git a/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
new file mode 100644
index 0000000..5e4bab15
--- /dev/null
+++ b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
@@ -0,0 +1,318 @@
+/*
+ * Copyright (C) 2017 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.server.backup;
+
+import static com.android.server.backup.BackupManagerService.MORE_DEBUG;
+import static com.android.server.backup.BackupManagerService.TAG;
+
+import android.annotation.Nullable;
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
+import android.app.ApplicationThreadConstants;
+import android.app.IActivityManager;
+import android.app.IBackupAgent;
+import android.app.backup.BackupAnnotations;
+import android.app.compat.CompatChanges;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledSince;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.content.pm.PackageManager.NameNotFoundException;
+import android.os.Binder;
+import android.os.Build;
+import android.os.IBinder;
+import android.os.Process;
+import android.os.RemoteException;
+import android.os.UserHandle;
+import android.util.ArraySet;
+import android.util.Slog;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.LocalServices;
+import com.android.server.backup.internal.LifecycleOperationStorage;
+
+import java.util.Set;
+
+/**
+ * Handles the lifecycle of {@link IBackupAgent}s that the {@link UserBackupManagerService}
+ * communicates with.
+ *
+ * <p>There can only be one agent that's connected to at a time.
+ *
+ * <p>There should be only one instance of this class per {@link UserBackupManagerService}.
+ */
+public class BackupAgentConnectionManager {
+
+    /**
+     * Enables the OS making a decision on whether backup restricted mode should be used for apps
+     * that haven't explicitly opted in or out. See
+     * {@link android.content.pm.PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE} for details.
+     */
+    @ChangeId
+    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.BAKLAVA)
+    public static final long OS_DECIDES_BACKUP_RESTRICTED_MODE = 376661510;
+
+    // The thread performing the sequence of queued backups binds to each app's agent
+    // in succession.  Bind notifications are asynchronously delivered through the
+    // Activity Manager; use this lock object to signal when a requested binding has
+    // completed.
+    private final Object mAgentConnectLock = new Object();
+    private IBackupAgent mConnectedAgent;
+    private volatile boolean mConnecting;
+    private final ArraySet<String> mRestoreNoRestrictedModePackages = new ArraySet<>();
+    private final ArraySet<String> mBackupNoRestrictedModePackages = new ArraySet<>();
+
+    private final IActivityManager mActivityManager;
+    private final ActivityManagerInternal mActivityManagerInternal;
+    private final LifecycleOperationStorage mOperationStorage;
+    private final PackageManager mPackageManager;
+    private final UserBackupManagerService mUserBackupManagerService;
+    private final int mUserId;
+    private final String mUserIdMsg;
+
+    BackupAgentConnectionManager(LifecycleOperationStorage operationStorage,
+            PackageManager packageManager, UserBackupManagerService userBackupManagerService,
+            int userId) {
+        mActivityManager = ActivityManager.getService();
+        mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
+        mOperationStorage = operationStorage;
+        mPackageManager = packageManager;
+        mUserBackupManagerService = userBackupManagerService;
+        mUserId = userId;
+        mUserIdMsg = "[UserID:" + userId + "] ";
+    }
+
+    /**
+     * Fires off a backup agent, blocking until it attaches (and ActivityManager will call
+     * {@link #agentConnected(String, IBinder)}) or until this operation times out.
+     *
+     * @param mode a {@code BACKUP_MODE} from {@link android.app.ApplicationThreadConstants}.
+     */
+    @Nullable
+    public IBackupAgent bindToAgentSynchronous(ApplicationInfo app, int mode,
+            @BackupAnnotations.BackupDestination int backupDestination) {
+        IBackupAgent agent = null;
+        synchronized (mAgentConnectLock) {
+            mConnecting = true;
+            mConnectedAgent = null;
+            boolean useRestrictedMode = shouldUseRestrictedBackupModeForPackage(mode,
+                    app.packageName);
+            try {
+                if (mActivityManager.bindBackupAgent(app.packageName, mode, mUserId,
+                        backupDestination, useRestrictedMode)) {
+                    Slog.d(TAG, mUserIdMsg + "awaiting agent for " + app);
+
+                    // success; wait for the agent to arrive
+                    // only wait 10 seconds for the bind to happen
+                    long timeoutMark = System.currentTimeMillis() + 10 * 1000;
+                    while (mConnecting && mConnectedAgent == null && (System.currentTimeMillis()
+                            < timeoutMark)) {
+                        try {
+                            mAgentConnectLock.wait(5000);
+                        } catch (InterruptedException e) {
+                            // just bail
+                            Slog.w(TAG, mUserIdMsg + "Interrupted: " + e);
+                            mConnecting = false;
+                            mConnectedAgent = null;
+                        }
+                    }
+
+                    // if we timed out with no connect, abort and move on
+                    if (mConnecting) {
+                        Slog.w(TAG, mUserIdMsg + "Timeout waiting for agent " + app);
+                        mConnectedAgent = null;
+                    }
+                    Slog.i(TAG, mUserIdMsg + "got agent " + mConnectedAgent);
+                    agent = mConnectedAgent;
+                }
+            } catch (RemoteException e) {
+                // can't happen - ActivityManager is local
+            }
+        }
+        if (agent == null) {
+            mActivityManagerInternal.clearPendingBackup(mUserId);
+        }
+        return agent;
+    }
+
+    /**
+     * Tell the ActivityManager that we are done with the {@link IBackupAgent} of this {@code app}.
+     * It will tell the app to destroy the agent.
+     */
+    public void unbindAgent(ApplicationInfo app) {
+        try {
+            mActivityManager.unbindBackupAgent(app);
+        } catch (RemoteException e) {
+            // Can't happen - activity manager is local
+        }
+    }
+
+    /**
+     * Callback: a requested backup agent has been instantiated. This should only be called from
+     * the
+     * {@link ActivityManager} when it's telling us that an agent is ready after a call to
+     * {@link #bindToAgentSynchronous(ApplicationInfo, int, int)}.
+     */
+    public void agentConnected(String packageName, IBinder agentBinder) {
+        synchronized (mAgentConnectLock) {
+            if (getCallingUid() == android.os.Process.SYSTEM_UID) {
+                Slog.d(TAG,
+                        mUserIdMsg + "agentConnected pkg=" + packageName + " agent=" + agentBinder);
+                mConnectedAgent = IBackupAgent.Stub.asInterface(agentBinder);
+                mConnecting = false;
+            } else {
+                Slog.w(TAG, mUserIdMsg + "Non-system process uid=" + getCallingUid()
+                        + " claiming agent connected");
+            }
+            mAgentConnectLock.notifyAll();
+        }
+    }
+
+    /**
+     * Callback: a backup agent has failed to come up, or has unexpectedly quit. If the agent failed
+     * to come up in the first place, the agentBinder argument will be {@code null}. This should
+     * only be called from the {@link ActivityManager}.
+     */
+    public void agentDisconnected(String packageName) {
+        synchronized (mAgentConnectLock) {
+            if (getCallingUid() == Process.SYSTEM_UID) {
+                mConnectedAgent = null;
+                mConnecting = false;
+            } else {
+                Slog.w(TAG, mUserIdMsg + "Non-system process uid=" + getCallingUid()
+                        + " claiming agent disconnected");
+            }
+            Slog.w(TAG, mUserIdMsg + "agentDisconnected: the backup agent for " + packageName
+                    + " died: cancel current operations");
+
+            // Offload operation cancellation off the main thread as the cancellation callbacks
+            // might call out to BackupTransport. Other operations started on the same package
+            // before the cancellation callback has executed will also be cancelled by the callback.
+            Runnable cancellationRunnable = () -> {
+                // handleCancel() causes the PerformFullTransportBackupTask to go on to
+                // tearDownAgentAndKill: that will unbindBackupAgent in the Activity Manager, so
+                // that the package being backed up doesn't get stuck in restricted mode until the
+                // backup time-out elapses.
+                for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
+                    if (MORE_DEBUG) {
+                        Slog.d(TAG,
+                                mUserIdMsg + "agentDisconnected: will handleCancel(all) for token:"
+                                        + Integer.toHexString(token));
+                    }
+                    mUserBackupManagerService.handleCancel(token, true /* cancelAll */);
+                }
+            };
+            getThreadForCancellation(cancellationRunnable).start();
+
+            mAgentConnectLock.notifyAll();
+        }
+    }
+
+    /**
+     * Marks the given set of packages as packages that should not be put into restricted mode if
+     * they are started for the given {@link BackupAnnotations.OperationType}.
+     */
+    public void setNoRestrictedModePackages(Set<String> packageNames,
+            @BackupAnnotations.OperationType int opType) {
+        if (opType == BackupAnnotations.OperationType.BACKUP) {
+            mBackupNoRestrictedModePackages.clear();
+            mBackupNoRestrictedModePackages.addAll(packageNames);
+        } else if (opType == BackupAnnotations.OperationType.RESTORE) {
+            mRestoreNoRestrictedModePackages.clear();
+            mRestoreNoRestrictedModePackages.addAll(packageNames);
+        } else {
+            throw new IllegalArgumentException("opType must be BACKUP or RESTORE");
+        }
+    }
+
+    /**
+     * Clears the list of packages that should not be put into restricted mode for either backup or
+     * restore.
+     */
+    public void clearNoRestrictedModePackages() {
+        mBackupNoRestrictedModePackages.clear();
+        mRestoreNoRestrictedModePackages.clear();
+    }
+
+    /**
+     * If the app has specified {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE}, then
+     * its value is returned. If it hasn't and it targets an SDK below
+     * {@link Build.VERSION_CODES#BAKLAVA} then returns true. If it targets a newer SDK, then
+     * returns the decision made by the {@link android.app.backup.BackupTransport}.
+     *
+     * <p>When this method is called, we should have already asked the transport and cached its
+     * response in {@link #mBackupNoRestrictedModePackages} or
+     * {@link #mRestoreNoRestrictedModePackages} so this method will immediately return without
+     * any IPC to the transport.
+     */
+    private boolean shouldUseRestrictedBackupModeForPackage(
+            @BackupAnnotations.OperationType int mode, String packageName) {
+        if (!Flags.enableRestrictedModeChanges()) {
+            return true;
+        }
+
+        // Key/Value apps are never put in restricted mode.
+        if (mode == ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL
+                || mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE) {
+            return false;
+        }
+
+        try {
+            PackageManager.Property property = mPackageManager.getPropertyAsUser(
+                    PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE,
+                    packageName, /* className= */ null, mUserId);
+            if (property.isBoolean()) {
+                // If the package has explicitly specified, we won't ask the transport.
+                return property.getBoolean();
+            } else {
+                Slog.w(TAG,
+                        PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE + "must be a boolean.");
+            }
+        } catch (NameNotFoundException e) {
+            // This is expected when the package has not defined the property in its manifest.
+        }
+
+        // The package has not specified the property. The behavior depends on the package's
+        // targetSdk.
+        // <36 gets the old behavior of always using restricted mode.
+        if (!CompatChanges.isChangeEnabled(OS_DECIDES_BACKUP_RESTRICTED_MODE, packageName,
+                UserHandle.of(mUserId))) {
+            return true;
+        }
+
+        // Apps targeting >=36 get the behavior decided by the transport.
+        // By this point, we should have asked the transport and cached its decision.
+        if ((mode == ApplicationThreadConstants.BACKUP_MODE_FULL
+                && mBackupNoRestrictedModePackages.contains(packageName)) || (
+                mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL
+                        && mRestoreNoRestrictedModePackages.contains(packageName))) {
+            Slog.d(TAG, "Transport requested no restricted mode for: " + packageName);
+            return false;
+        }
+        return true;
+    }
+
+    @VisibleForTesting
+    Thread getThreadForCancellation(Runnable operation) {
+        return new Thread(operation, /* operationName */ "agent-disconnected");
+    }
+
+    @VisibleForTesting
+    int getCallingUid() {
+        return Binder.getCallingUid();
+    }
+}
diff --git a/services/backup/java/com/android/server/backup/BackupManagerService.java b/services/backup/java/com/android/server/backup/BackupManagerService.java
index 5f0071d..3f6ede9 100644
--- a/services/backup/java/com/android/server/backup/BackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/BackupManagerService.java
@@ -658,7 +658,8 @@
                 getServiceForUserIfCallerHasPermission(userId, "agentConnected()");
 
         if (userBackupManagerService != null) {
-            userBackupManagerService.agentConnected(packageName, agentBinder);
+            userBackupManagerService.getBackupAgentConnectionManager().agentConnected(packageName,
+                    agentBinder);
         }
     }
 
@@ -683,7 +684,8 @@
                 getServiceForUserIfCallerHasPermission(userId, "agentDisconnected()");
 
         if (userBackupManagerService != null) {
-            userBackupManagerService.agentDisconnected(packageName);
+            userBackupManagerService.getBackupAgentConnectionManager().agentDisconnected(
+                    packageName);
         }
     }
 
diff --git a/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java b/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
index f5d6836..136bacd 100644
--- a/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
+++ b/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
@@ -146,7 +146,8 @@
 
     private IBackupAgent bindToAgent(ApplicationInfo targetApp) {
         try {
-            return mBackupManagerService.bindToAgentSynchronous(targetApp,
+            return mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
+                    targetApp,
                     ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL,
                     BackupAnnotations.BackupDestination.CLOUD);
         } catch (SecurityException e) {
diff --git a/services/backup/java/com/android/server/backup/UserBackupManagerService.java b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
index 5de2fb3..e085f6e 100644
--- a/services/backup/java/com/android/server/backup/UserBackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
@@ -43,9 +43,7 @@
 import android.app.ActivityManagerInternal;
 import android.app.AlarmManager;
 import android.app.AppGlobals;
-import android.app.ApplicationThreadConstants;
 import android.app.IActivityManager;
-import android.app.IBackupAgent;
 import android.app.PendingIntent;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupAnnotations;
@@ -60,9 +58,6 @@
 import android.app.backup.IFullBackupRestoreObserver;
 import android.app.backup.IRestoreSession;
 import android.app.backup.ISelectBackupTransportCallback;
-import android.app.compat.CompatChanges;
-import android.compat.annotation.ChangeId;
-import android.compat.annotation.EnabledSince;
 import android.content.ActivityNotFoundException;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
@@ -83,7 +78,6 @@
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.HandlerThread;
-import android.os.IBinder;
 import android.os.Message;
 import android.os.ParcelFileDescriptor;
 import android.os.PowerManager;
@@ -302,21 +296,10 @@
     private static final String BACKUP_FINISHED_ACTION = "android.intent.action.BACKUP_FINISHED";
     private static final String BACKUP_FINISHED_PACKAGE_EXTRA = "packageName";
 
-    /**
-     * Enables the OS making a decision on whether backup restricted mode should be used for apps
-     * that haven't explicitly opted in or out. See
-     * {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE} for details.
-     */
-    @ChangeId
-    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.BAKLAVA)
-    public static final long OS_DECIDES_BACKUP_RESTRICTED_MODE = 376661510;
-
     // Time delay for initialization operations that can be delayed so as not to consume too much
     // CPU on bring-up and increase time-to-UI.
     private static final long INITIALIZATION_DELAY_MILLIS = 3000;
 
-    // Timeout interval for deciding that a bind has taken too long.
-    private static final long BIND_TIMEOUT_INTERVAL = 10 * 1000;
     // Timeout interval for deciding that a clear-data has taken too long.
     private static final long CLEAR_DATA_TIMEOUT_INTERVAL = 30 * 1000;
 
@@ -365,22 +348,9 @@
     // Backups that we haven't started yet.  Keys are package names.
     private final HashMap<String, BackupRequest> mPendingBackups = new HashMap<>();
 
-    private final ArraySet<String> mRestoreNoRestrictedModePackages = new ArraySet<>();
-    private final ArraySet<String> mBackupNoRestrictedModePackages = new ArraySet<>();
-
     // locking around the pending-backup management
     private final Object mQueueLock = new Object();
-
     private final UserBackupPreferences mBackupPreferences;
-
-    // The thread performing the sequence of queued backups binds to each app's agent
-    // in succession.  Bind notifications are asynchronously delivered through the
-    // Activity Manager; use this lock object to signal when a requested binding has
-    // completed.
-    private final Object mAgentConnectLock = new Object();
-    private IBackupAgent mConnectedAgent;
-    private volatile boolean mConnecting;
-
     private volatile boolean mBackupRunning;
     private volatile long mLastBackupPass;
 
@@ -410,6 +380,7 @@
 
     private ActiveRestoreSession mActiveRestoreSession;
 
+    private final BackupAgentConnectionManager mBackupAgentConnectionManager;
     private final LifecycleOperationStorage mOperationStorage;
 
     private final Random mTokenGenerator = new Random();
@@ -547,6 +518,8 @@
         mRegisterTransportsRequestedTime = 0;
         mPackageManager = packageManager;
         mOperationStorage = operationStorage;
+        mBackupAgentConnectionManager = new BackupAgentConnectionManager(mOperationStorage,
+                mPackageManager, this, mUserId);
         mTransportManager = transportManager;
         mFullBackupQueue = new ArrayList<>();
         mBackupHandler = backupHandler;
@@ -599,6 +572,8 @@
         mAgentTimeoutParameters.start();
 
         mOperationStorage = new LifecycleOperationStorage(mUserId);
+        mBackupAgentConnectionManager = new BackupAgentConnectionManager(mOperationStorage,
+                mPackageManager, this, mUserId);
 
         Objects.requireNonNull(userBackupThread, "userBackupThread cannot be null");
         mBackupHandler = new BackupHandler(this, mOperationStorage, userBackupThread);
@@ -1660,67 +1635,6 @@
         }
     }
 
-    /** Fires off a backup agent, blocking until it attaches or times out. */
-    @Nullable
-    public IBackupAgent bindToAgentSynchronous(ApplicationInfo app, int mode,
-            @BackupDestination int backupDestination) {
-        IBackupAgent agent = null;
-        synchronized (mAgentConnectLock) {
-            mConnecting = true;
-            mConnectedAgent = null;
-            boolean useRestrictedMode = shouldUseRestrictedBackupModeForPackage(mode,
-                    app.packageName);
-            try {
-                if (mActivityManager.bindBackupAgent(app.packageName, mode, mUserId,
-                        backupDestination, useRestrictedMode)) {
-                    Slog.d(TAG, addUserIdToLogMessage(mUserId, "awaiting agent for " + app));
-
-                    // success; wait for the agent to arrive
-                    // only wait 10 seconds for the bind to happen
-                    long timeoutMark = System.currentTimeMillis() + BIND_TIMEOUT_INTERVAL;
-                    while (mConnecting && mConnectedAgent == null
-                            && (System.currentTimeMillis() < timeoutMark)) {
-                        try {
-                            mAgentConnectLock.wait(5000);
-                        } catch (InterruptedException e) {
-                            // just bail
-                            Slog.w(TAG, addUserIdToLogMessage(mUserId, "Interrupted: " + e));
-                            mConnecting = false;
-                            mConnectedAgent = null;
-                        }
-                    }
-
-                    // if we timed out with no connect, abort and move on
-                    if (mConnecting) {
-                        Slog.w(
-                                TAG,
-                                addUserIdToLogMessage(mUserId, "Timeout waiting for agent " + app));
-                        mConnectedAgent = null;
-                    }
-                    if (DEBUG) {
-                        Slog.i(TAG, addUserIdToLogMessage(mUserId, "got agent " + mConnectedAgent));
-                    }
-                    agent = mConnectedAgent;
-                }
-            } catch (RemoteException e) {
-                // can't happen - ActivityManager is local
-            }
-        }
-        if (agent == null) {
-            mActivityManagerInternal.clearPendingBackup(mUserId);
-        }
-        return agent;
-    }
-
-    /** Unbind from a backup agent. */
-    public void unbindAgent(ApplicationInfo app) {
-        try {
-            mActivityManager.unbindBackupAgent(app);
-        } catch (RemoteException e) {
-            // Can't happen - activity manager is local
-        }
-    }
-
     /**
      * Clear an application's data after a failed restore, blocking until the operation completes or
      * times out.
@@ -2493,10 +2407,6 @@
         AppWidgetBackupBridge.restoreWidgetState(packageName, widgetData, mUserId);
     }
 
-    // *****************************
-    // NEW UNIFIED RESTORE IMPLEMENTATION
-    // *****************************
-
     /** Schedule a backup pass for {@code packageName}. */
     public void dataChangedImpl(String packageName) {
         HashSet<String> targets = dataChangedTargets(packageName);
@@ -3122,91 +3032,6 @@
         }
     }
 
-    /**
-     * Marks the given set of packages as packages that should not be put into restricted mode if
-     * they are started for the given {@link BackupAnnotations.OperationType}.
-     */
-    public void setNoRestrictedModePackages(Set<String> packageNames,
-            @BackupAnnotations.OperationType int opType) {
-        if (opType == BackupAnnotations.OperationType.BACKUP) {
-            mBackupNoRestrictedModePackages.clear();
-            mBackupNoRestrictedModePackages.addAll(packageNames);
-        } else if (opType == BackupAnnotations.OperationType.RESTORE) {
-            mRestoreNoRestrictedModePackages.clear();
-            mRestoreNoRestrictedModePackages.addAll(packageNames);
-        } else {
-            throw new IllegalArgumentException("opType must be BACKUP or RESTORE");
-        }
-    }
-
-    /**
-     * Clears the list of packages that should not be put into restricted mode for either backup or
-     * restore.
-     */
-    public void clearNoRestrictedModePackages() {
-        mBackupNoRestrictedModePackages.clear();
-        mRestoreNoRestrictedModePackages.clear();
-    }
-
-    /**
-     * If the app has specified {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE}, then
-     * its value is returned. If it hasn't and it targets an SDK below
-     * {@link Build.VERSION_CODES#BAKLAVA} then returns true. If it targets a newer SDK, then
-     * returns the decision made by the {@link android.app.backup.BackupTransport}.
-     *
-     * <p>When this method is called, we should have already asked the transport and cached its
-     * response in {@link #mBackupNoRestrictedModePackages} or
-     * {@link #mRestoreNoRestrictedModePackages} so this method will immediately return without
-     * any IPC to the transport.
-     */
-    private boolean shouldUseRestrictedBackupModeForPackage(
-            @BackupAnnotations.OperationType int mode, String packageName) {
-        if (!Flags.enableRestrictedModeChanges()) {
-            return true;
-        }
-
-        // Key/Value apps are never put in restricted mode.
-        if (mode == ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL
-                || mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE) {
-            return false;
-        }
-
-        try {
-            PackageManager.Property property = mPackageManager.getPropertyAsUser(
-                    PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE,
-                    packageName, /* className= */ null,
-                    mUserId);
-            if (property.isBoolean()) {
-                // If the package has explicitly specified, we won't ask the transport.
-                return property.getBoolean();
-            } else {
-                Slog.w(TAG, PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE
-                        + "must be a boolean.");
-            }
-        } catch (NameNotFoundException e) {
-            // This is expected when the package has not defined the property in its manifest.
-        }
-
-        // The package has not specified the property. The behavior depends on the package's
-        // targetSdk.
-        // <36 gets the old behavior of always using restricted mode.
-        if (!CompatChanges.isChangeEnabled(OS_DECIDES_BACKUP_RESTRICTED_MODE, packageName,
-                UserHandle.of(mUserId))) {
-            return true;
-        }
-
-        // Apps targeting >=36 get the behavior decided by the transport.
-        // By this point, we should have asked the transport and cached its decision.
-        if ((mode == ApplicationThreadConstants.BACKUP_MODE_FULL
-                && mBackupNoRestrictedModePackages.contains(packageName))
-                || (mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL
-                && mRestoreNoRestrictedModePackages.contains(packageName))) {
-            Slog.d(TAG, "Transport requested no restricted mode for: " + packageName);
-            return false;
-        }
-        return true;
-    }
-
     private boolean startConfirmationUi(int token, String action) {
         try {
             Intent confIntent = new Intent(action);
@@ -3898,83 +3723,6 @@
     }
 
     /**
-     * Callback: a requested backup agent has been instantiated. This should only be called from the
-     * {@link ActivityManager}.
-     */
-    public void agentConnected(String packageName, IBinder agentBinder) {
-        synchronized (mAgentConnectLock) {
-            if (Binder.getCallingUid() == Process.SYSTEM_UID) {
-                Slog.d(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "agentConnected pkg=" + packageName + " agent=" + agentBinder));
-                mConnectedAgent = IBackupAgent.Stub.asInterface(agentBinder);
-                mConnecting = false;
-            } else {
-                Slog.w(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "Non-system process uid="
-                                        + Binder.getCallingUid()
-                                        + " claiming agent connected"));
-            }
-            mAgentConnectLock.notifyAll();
-        }
-    }
-
-    /**
-     * Callback: a backup agent has failed to come up, or has unexpectedly quit. If the agent failed
-     * to come up in the first place, the agentBinder argument will be {@code null}. This should
-     * only be called from the {@link ActivityManager}.
-     */
-    public void agentDisconnected(String packageName) {
-        synchronized (mAgentConnectLock) {
-            if (Binder.getCallingUid() == Process.SYSTEM_UID) {
-                mConnectedAgent = null;
-                mConnecting = false;
-            } else {
-                Slog.w(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "Non-system process uid="
-                                        + Binder.getCallingUid()
-                                        + " claiming agent disconnected"));
-            }
-            Slog.w(TAG, "agentDisconnected: the backup agent for " + packageName
-                    + " died: cancel current operations");
-
-            // Offload operation cancellation off the main thread as the cancellation callbacks
-            // might call out to BackupTransport. Other operations started on the same package
-            // before the cancellation callback has executed will also be cancelled by the callback.
-            Runnable cancellationRunnable = () -> {
-                // handleCancel() causes the PerformFullTransportBackupTask to go on to
-                // tearDownAgentAndKill: that will unbindBackupAgent in the Activity Manager, so
-                // that the package being backed up doesn't get stuck in restricted mode until the
-                // backup time-out elapses.
-                for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
-                    if (MORE_DEBUG) {
-                        Slog.d(TAG, "agentDisconnected: will handleCancel(all) for token:"
-                                + Integer.toHexString(token));
-                    }
-                    handleCancel(token, true /* cancelAll */);
-                }
-            };
-            getThreadForAsyncOperation(/* operationName */ "agent-disconnected",
-                    cancellationRunnable).start();
-
-            mAgentConnectLock.notifyAll();
-        }
-    }
-
-    @VisibleForTesting
-    Thread getThreadForAsyncOperation(String operationName, Runnable operation) {
-        return new Thread(operation, operationName);
-    }
-
-    /**
      * An application being installed will need a restore pass, then the {@link PackageManager} will
      * need to be told when the restore is finished.
      */
@@ -4521,4 +4269,8 @@
     public IBackupManager getBackupManagerBinder() {
         return mBackupManagerBinder;
     }
+
+    public BackupAgentConnectionManager getBackupAgentConnectionManager() {
+        return mBackupAgentConnectionManager;
+    }
 }
diff --git a/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java b/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
index 1271206..b98cb10 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
@@ -314,7 +314,7 @@
                 Slog.d(TAG, "Binding to full backup agent : " + mPkg.packageName);
             }
             mAgent =
-                    backupManagerService.bindToAgentSynchronous(
+                    backupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                             mPkg.applicationInfo, ApplicationThreadConstants.BACKUP_MODE_FULL,
                             mBackupEligibilityRules.getBackupDestination());
         }
diff --git a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
index be9cdc8..65730c9 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
@@ -702,7 +702,8 @@
             }
 
             // Clear this to avoid using the memory until reboot.
-            mUserBackupManagerService.clearNoRestrictedModePackages();
+            mUserBackupManagerService
+                    .getBackupAgentConnectionManager().clearNoRestrictedModePackages();
 
             Slog.i(TAG, "Full data backup pass finished.");
             mUserBackupManagerService.getWakelock().release();
@@ -741,7 +742,8 @@
             }
             packageNames = transport.getPackagesThatShouldNotUseRestrictedMode(packageNames,
                     BACKUP);
-            mUserBackupManagerService.setNoRestrictedModePackages(packageNames, BACKUP);
+            mUserBackupManagerService.getBackupAgentConnectionManager().setNoRestrictedModePackages(
+                    packageNames, BACKUP);
         } catch (RemoteException e) {
             Slog.i(TAG, "Failed to retrieve no restricted mode packages from transport");
         }
diff --git a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
index 3a6e1ca..82232a6 100644
--- a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
+++ b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
@@ -741,7 +741,7 @@
         final IBackupAgent agent;
         try {
             agent =
-                    mBackupManagerService.bindToAgentSynchronous(
+                    mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                             packageInfo.applicationInfo, BACKUP_MODE_INCREMENTAL,
                             mBackupEligibilityRules.getBackupDestination());
             if (agent == null) {
@@ -1302,7 +1302,8 @@
 
         // For PM metadata (for which applicationInfo is null) there is no agent-bound state.
         if (mCurrentPackage.applicationInfo != null) {
-            mBackupManagerService.unbindAgent(mCurrentPackage.applicationInfo);
+            mBackupManagerService.getBackupAgentConnectionManager().unbindAgent(
+                    mCurrentPackage.applicationInfo);
         }
         mAgent = null;
     }
diff --git a/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java b/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
index 2d99c96..b59e860 100644
--- a/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
+++ b/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
@@ -410,11 +410,7 @@
 
                             // All set; now set up the IPC and launch the agent
                             setUpPipes();
-                            mAgent = mBackupManagerService.bindToAgentSynchronous(mTargetApp,
-                                    FullBackup.KEY_VALUE_DATA_TOKEN.equals(info.domain)
-                                            ? ApplicationThreadConstants.BACKUP_MODE_RESTORE
-                                            : ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL,
-                                    mBackupEligibilityRules.getBackupDestination());
+                            mAgent = bindToAgent(info);
                             mAgentPackage = pkg;
                         } catch (IOException | NameNotFoundException e) {
                             // fall through to error handling
@@ -805,15 +801,12 @@
         return packages.contains(packageName);
     }
 
-    void sendOnRestorePackage(String name) {
-        if (mObserver != null) {
-            try {
-                // TODO: use a more user-friendly name string
-                mObserver.onRestorePackage(name);
-            } catch (RemoteException e) {
-                Slog.w(TAG, "full restore observer went away: restorePackage");
-                mObserver = null;
-            }
-        }
+    private IBackupAgent bindToAgent(FileMetadata info) {
+        return mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
+                mTargetApp,
+                FullBackup.KEY_VALUE_DATA_TOKEN.equals(info.domain)
+                        ? ApplicationThreadConstants.BACKUP_MODE_RESTORE
+                        : ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL,
+                mBackupEligibilityRules.getBackupDestination());
     }
 }
diff --git a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
index 5ee51a5..e5c7e5c 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -814,7 +814,7 @@
 
         // Good to go!  Set up and bind the agent...
         mAgent =
-                backupManagerService.bindToAgentSynchronous(
+                backupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                         mCurrentPackage.applicationInfo,
                         ApplicationThreadConstants.BACKUP_MODE_RESTORE,
                         mBackupEligibilityRules.getBackupDestination());
@@ -1364,7 +1364,7 @@
         backupManagerService.getBackupHandler().removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
 
         // Clear this to avoid using the memory until reboot.
-        backupManagerService.clearNoRestrictedModePackages();
+        backupManagerService.getBackupAgentConnectionManager().clearNoRestrictedModePackages();
 
         // If we have a PM token, we must under all circumstances be sure to
         // handshake when we've finished.
@@ -1838,7 +1838,8 @@
             }
             packageNames = transport.getPackagesThatShouldNotUseRestrictedMode(packageNames,
                     RESTORE);
-            backupManagerService.setNoRestrictedModePackages(packageNames, RESTORE);
+            backupManagerService.getBackupAgentConnectionManager().setNoRestrictedModePackages(
+                    packageNames, RESTORE);
         } catch (RemoteException e) {
             Slog.i(TAG, "Failed to retrieve restricted mode packages from transport");
         }
diff --git a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
index a547d0f..4e9fff2 100644
--- a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
@@ -31,6 +31,7 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.robolectric.Shadows.shadowOf;
 import static org.testng.Assert.expectThrows;
 
@@ -96,6 +97,7 @@
     @UserIdInt private int mUserTwoId;
     @Mock private UserBackupManagerService mUserSystemService;
     @Mock private UserBackupManagerService mUserOneService;
+    @Mock private BackupAgentConnectionManager mUserOneBackupAgentConnectionManager;
     @Mock private UserBackupManagerService mUserTwoService;
 
     /** Setup */
@@ -116,6 +118,9 @@
         mShadowContext.grantPermissions(BACKUP);
         mShadowContext.grantPermissions(INTERACT_ACROSS_USERS_FULL);
 
+        when(mUserOneService.getBackupAgentConnectionManager()).thenReturn(
+                mUserOneBackupAgentConnectionManager);
+
         ShadowBinder.setCallingUid(Process.SYSTEM_UID);
     }
 
@@ -226,7 +231,7 @@
 
         backupManagerService.agentConnected(mUserOneId, TEST_PACKAGE, agentBinder);
 
-        verify(mUserOneService).agentConnected(TEST_PACKAGE, agentBinder);
+        verify(mUserOneBackupAgentConnectionManager).agentConnected(TEST_PACKAGE, agentBinder);
     }
 
     /** Test that the backup service does not route methods for non-registered users. */
@@ -239,7 +244,8 @@
 
         backupManagerService.agentConnected(mUserTwoId, TEST_PACKAGE, agentBinder);
 
-        verify(mUserOneService, never()).agentConnected(TEST_PACKAGE, agentBinder);
+        verify(mUserOneBackupAgentConnectionManager, never()).agentConnected(TEST_PACKAGE,
+                agentBinder);
     }
 
     /** Test that the backup service routes methods correctly to the user that requests it. */
diff --git a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
index 7349c14..aeb1ba9 100644
--- a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
@@ -107,6 +107,7 @@
 import com.android.internal.infra.AndroidFuture;
 import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.BackupRestoreTask;
 import com.android.server.backup.DataChangedJournal;
 import com.android.server.backup.KeyValueBackupJob;
@@ -167,7 +168,6 @@
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Stream;
 
-// TODO: Test agents timing out
 @RunWith(RobolectricTestRunner.class)
 @Config(
         shadows = {
@@ -195,6 +195,7 @@
     @Mock private IBackupManagerMonitor mMonitor;
     @Mock private OnTaskFinishedListener mListener;
     @Mock private PackageManagerInternal mPackageManagerInternal;
+    @Mock private BackupAgentConnectionManager mBackupAgentConnectionManager;
 
     private UserBackupManagerService mBackupManagerService;
     private TransportData mTransport;
@@ -257,6 +258,8 @@
         when(mBackupManagerService.getBaseStateDir()).thenReturn(mBaseStateDir);
         when(mBackupManagerService.getDataDir()).thenReturn(mDataDir);
         when(mBackupManagerService.getBackupManagerBinder()).thenReturn(mBackupManager);
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
 
         mBackupHandler = mBackupManagerService.getBackupHandler();
         mShadowBackupLooper = shadowOf(mBackupHandler.getLooper());
@@ -749,7 +752,8 @@
 
     /**
      * Agent unavailable means {@link
-     * UserBackupManagerService#bindToAgentSynchronous(ApplicationInfo, int)} returns {@code null}.
+     * BackupAgentConnectionManager#bindToAgentSynchronous(ApplicationInfo, int, int)} returns
+     * {@code null}.
      *
      * @see #setUpAgent(PackageData)
      */
@@ -805,7 +809,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgent(PACKAGE_1);
         doThrow(SecurityException.class)
-                .when(mBackupManagerService)
+                .when(mBackupAgentConnectionManager)
                 .bindToAgentSynchronous(argThat(applicationInfo(PACKAGE_1)), anyInt(), anyInt());
         KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
@@ -823,7 +827,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgent(PACKAGE_1);
         doThrow(SecurityException.class)
-                .when(mBackupManagerService)
+                .when(mBackupAgentConnectionManager)
                 .bindToAgentSynchronous(argThat(applicationInfo(PACKAGE_1)), anyInt(), anyInt());
         KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true, PACKAGE_1);
 
@@ -861,7 +865,7 @@
         runTask(task);
 
         verify(mBackupManagerService).setWorkSource(null);
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -1097,7 +1101,7 @@
         runTask(task);
 
         verify(agentMock.agentBinder).fail(any());
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -1418,7 +1422,8 @@
                 .isEqualTo("newState".getBytes());
         assertCleansUpFiles(mTransport, PM_PACKAGE);
         // We don't unbind PM
-        verify(mBackupManagerService, never()).unbindAgent(argThat(applicationInfo(PM_PACKAGE)));
+        verify(mBackupAgentConnectionManager, never()).unbindAgent(
+                argThat(applicationInfo(PM_PACKAGE)));
     }
 
     @Test
@@ -1439,7 +1444,8 @@
 
         runTask(task);
 
-        verify(mBackupManagerService, never()).unbindAgent(argThat(applicationInfo(PM_PACKAGE)));
+        verify(mBackupAgentConnectionManager, never()).unbindAgent(
+                argThat(applicationInfo(PM_PACKAGE)));
     }
 
     @Test
@@ -1642,9 +1648,10 @@
 
         runTask(task);
 
-        InOrder inOrder = inOrder(agentMock.agent, mBackupManagerService);
+        InOrder inOrder = inOrder(agentMock.agent, mBackupAgentConnectionManager);
         inOrder.verify(agentMock.agent).onQuotaExceeded(anyLong(), eq(1234L));
-        inOrder.verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        inOrder.verify(mBackupAgentConnectionManager).unbindAgent(
+                argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -2634,12 +2641,12 @@
             doNothing().when(backupAgentBinder).fail(any());
             if (packageData.available) {
                 doReturn(backupAgentBinder)
-                        .when(mBackupManagerService)
+                        .when(mBackupAgentConnectionManager)
                         .bindToAgentSynchronous(argThat(applicationInfo(packageData)), anyInt(),
                                 anyInt());
             } else {
                 doReturn(null)
-                        .when(mBackupManagerService)
+                        .when(mBackupAgentConnectionManager)
                         .bindToAgentSynchronous(argThat(applicationInfo(packageData)), anyInt(),
                                 anyInt());
             }
@@ -2976,7 +2983,7 @@
 
     private void assertCleansUpFilesAndAgent(TransportData transport, PackageData packageData) {
         assertCleansUpFiles(transport, packageData);
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(packageData)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(packageData)));
     }
 
     private void assertCleansUpFiles(TransportData transport, PackageData packageData) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
new file mode 100644
index 0000000..19e43b6
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
@@ -0,0 +1,381 @@
+/*
+ * 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.server.backup;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
+import android.app.ApplicationThreadConstants;
+import android.app.IActivityManager;
+import android.app.IBackupAgent;
+import android.app.backup.BackupAnnotations.BackupDestination;
+import android.app.backup.BackupAnnotations.OperationType;
+import android.compat.testing.PlatformCompatChangeRule;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.os.Process;
+import android.os.UserHandle;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.annotations.Presubmit;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.server.LocalServices;
+import com.android.server.backup.internal.LifecycleOperationStorage;
+
+import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
+import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
+
+import com.google.common.collect.ImmutableSet;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+import java.util.Set;
+
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class BackupAgentConnectionManagerTest {
+    private static final String TEST_PACKAGE = "com.test.package";
+
+    @Rule
+    public TestRule compatChangeRule = new PlatformCompatChangeRule();
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    @Mock
+    IActivityManager mActivityManager;
+    @Mock
+    ActivityManagerInternal mActivityManagerInternal;
+    @Mock
+    LifecycleOperationStorage mOperationStorage;
+    @Mock
+    UserBackupManagerService mUserBackupManagerService;
+    @Mock
+    IBackupAgent.Stub mBackupAgentStub;
+    @Mock
+    PackageManager mPackageManager;
+
+    private BackupAgentConnectionManager mConnectionManager;
+    private MockitoSession mSession;
+    private ApplicationInfo mTestApplicationInfo;
+    private IBackupAgent mBackupAgentResult;
+    private Thread mTestThread;
+
+    @Before
+    public void setUp() throws Exception {
+        mSession = mockitoSession().initMocks(this).mockStatic(ActivityManager.class).mockStatic(
+                LocalServices.class).strictness(Strictness.LENIENT).startMocking();
+        MockitoAnnotations.initMocks(this);
+
+        doReturn(mActivityManager).when(ActivityManager::getService);
+        doReturn(mActivityManagerInternal).when(
+                () -> LocalServices.getService(ActivityManagerInternal.class));
+        // Real package manager throws if a property is not defined.
+        when(mPackageManager.getPropertyAsUser(any(), any(), any(), anyInt())).thenThrow(
+                new PackageManager.NameNotFoundException());
+
+        mConnectionManager = spy(
+                new BackupAgentConnectionManager(mOperationStorage, mPackageManager,
+                        mUserBackupManagerService, UserHandle.USER_SYSTEM));
+
+        mTestApplicationInfo = new ApplicationInfo();
+        mTestApplicationInfo.packageName = TEST_PACKAGE;
+
+        mBackupAgentResult = null;
+        mTestThread = null;
+    }
+
+    @After
+    public void tearDown() {
+        if (mSession != null) {
+            mSession.finishMocking();
+        }
+    }
+
+    @Test
+    public void bindToAgentSynchronous_amReturnsFailure_returnsNullAndClearsPendingBackups()
+            throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(false);
+
+        IBackupAgent result = mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        assertThat(result).isNull();
+        verify(mActivityManagerInternal).clearPendingBackup(UserHandle.USER_SYSTEM);
+    }
+
+    @Test
+    public void bindToAgentSynchronous_agentDisconnectedCalled_returnsNullAndClearsPendingBackups()
+            throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(true);
+        // This is so that IBackupAgent.Stub.asInterface() works.
+        when(mBackupAgentStub.queryLocalInterface(any())).thenReturn(mBackupAgentStub);
+        when(mConnectionManager.getCallingUid()).thenReturn(Process.SYSTEM_UID);
+
+        // This is going to block until it receives the callback so we need to run it on a
+        // separate thread.
+        Thread testThread = new Thread(() -> setBackupAgentResult(
+                mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                        ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD)),
+                "backup-agent-connection-manager-test");
+        testThread.start();
+        // Give the testThread a head start, otherwise agentConnected() might run before
+        // bindToAgentSynchronous() is called.
+        Thread.sleep(500);
+        mConnectionManager.agentDisconnected(TEST_PACKAGE);
+        testThread.join();
+
+        assertThat(mBackupAgentResult).isNull();
+        verify(mActivityManagerInternal).clearPendingBackup(UserHandle.USER_SYSTEM);
+    }
+
+    @Test
+    public void bindToAgentSynchronous_agentConnectedCalled_returnsBackupAgent() throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(true);
+        // This is so that IBackupAgent.Stub.asInterface() works.
+        when(mBackupAgentStub.queryLocalInterface(any())).thenReturn(mBackupAgentStub);
+        when(mConnectionManager.getCallingUid()).thenReturn(Process.SYSTEM_UID);
+
+        // This is going to block until it receives the callback so we need to run it on a
+        // separate thread.
+        Thread testThread = new Thread(() -> setBackupAgentResult(
+                mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                        ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD)),
+                "backup-agent-connection-manager-test");
+        testThread.start();
+        // Give the testThread a head start, otherwise agentConnected() might run before
+        // bindToAgentSynchronous() is called.
+        Thread.sleep(500);
+        mConnectionManager.agentConnected(TEST_PACKAGE, mBackupAgentStub);
+        testThread.join();
+
+        assertThat(mBackupAgentResult).isEqualTo(mBackupAgentStub);
+        verify(mActivityManagerInternal, never()).clearPendingBackup(anyInt());
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_restrictedModeChangesFlagOff_shouldUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_keyValueBackup_shouldNotUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_keyValueRestore_shouldNotUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_RESTORE, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_packageOptedIn_shouldUseRestrictedMode() throws Exception {
+        reset(mPackageManager);
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenReturn(new PackageManager.Property(
+                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ true,
+                TEST_PACKAGE, /* className= */ null));
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_packageOptedOut_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenReturn(new PackageManager.Property(
+                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ false,
+                TEST_PACKAGE, /* className= */ null));
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @DisableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_targetSdkBelowB_shouldUseRestrictedMode() throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_targetSdkB_notInList_shouldUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.clearNoRestrictedModePackages();
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_forRestore_targetSdkB_inList_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.RESTORE);
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_forBackup_targetSdkB_inList_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.BACKUP);
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    public void agentDisconnected_cancelsCurrentOperations() throws Exception {
+        when(mOperationStorage.operationTokensForPackage(eq(TEST_PACKAGE))).thenReturn(
+                ImmutableSet.of(123, 456, 789));
+        when(mConnectionManager.getThreadForCancellation(any())).thenAnswer(invocation -> {
+            Thread testThread = new Thread((Runnable) invocation.getArgument(0),
+                    "agent-disconnected-test");
+            setTestThread(testThread);
+            return testThread;
+        });
+
+        mConnectionManager.agentDisconnected(TEST_PACKAGE);
+
+        mTestThread.join();
+        verify(mUserBackupManagerService).handleCancel(eq(123), eq(true));
+        verify(mUserBackupManagerService).handleCancel(eq(456), eq(true));
+        verify(mUserBackupManagerService).handleCancel(eq(789), eq(true));
+    }
+
+    @Test
+    public void unbindAgent_callsAmUnbindBackupAgent() throws Exception {
+        mConnectionManager.unbindAgent(mTestApplicationInfo);
+
+        verify(mActivityManager).unbindBackupAgent(eq(mTestApplicationInfo));
+    }
+
+    // Needed because variables can't be assigned directly inside lambdas in Java.
+    private void setBackupAgentResult(IBackupAgent result) {
+        mBackupAgentResult = result;
+    }
+
+    // Needed because variables can't be assigned directly inside lambdas in Java.
+    private void setTestThread(Thread thread) {
+        mTestThread = thread;
+    }
+}
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
index 07f2188..7a9e96f 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
@@ -18,20 +18,18 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.fail;
+
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.when;
 
 import android.annotation.UserIdInt;
 import android.app.ActivityManagerInternal;
-import android.app.ApplicationThreadConstants;
 import android.app.IActivityManager;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupAnnotations.BackupDestination;
@@ -47,8 +45,6 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.os.Handler;
-import android.platform.test.annotations.DisableFlags;
-import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
@@ -57,7 +53,6 @@
 import android.util.KeyValueListParser;
 
 import androidx.test.core.app.ApplicationProvider;
-import androidx.test.filters.FlakyTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.server.backup.internal.BackupHandler;
@@ -69,8 +64,6 @@
 import com.android.server.backup.utils.BackupEligibilityRules;
 import com.android.server.backup.utils.BackupManagerMonitorEventSender;
 
-import com.google.common.collect.ImmutableSet;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -84,11 +77,6 @@
 
 import java.util.Arrays;
 import java.util.List;
-import java.util.Set;
-import java.util.function.IntConsumer;
-
-import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
-import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 @Presubmit
 @RunWith(AndroidJUnit4.class)
@@ -96,7 +84,6 @@
     private static final String TEST_PACKAGE = "package1";
     private static final String[] TEST_PACKAGES = new String[] { TEST_PACKAGE };
     private static final String TEST_TRANSPORT = "transport";
-    private static final int WORKER_THREAD_TIMEOUT_MILLISECONDS = 100;
     @UserIdInt private static final int USER_ID = 0;
 
     @Rule
@@ -278,33 +265,6 @@
     }
 
     @Test
-    @FlakyTest
-    public void testAgentDisconnected_cancelsCurrentOperations() throws Exception {
-        when(mOperationStorage.operationTokensForPackage(eq("com.android.foo"))).thenReturn(
-                ImmutableSet.of(123, 456, 789)
-        );
-
-        mService.agentDisconnected("com.android.foo");
-
-        mService.waitForAsyncOperation();
-        verify(mOperationStorage).cancelOperation(eq(123), eq(true), any(IntConsumer.class));
-        verify(mOperationStorage).cancelOperation(eq(456), eq(true), any());
-        verify(mOperationStorage).cancelOperation(eq(789), eq(true), any());
-    }
-
-    @Test
-    public void testAgentDisconnected_unknownPackageName_cancelsNothing() throws Exception {
-        when(mOperationStorage.operationTokensForPackage(eq("com.android.foo"))).thenReturn(
-                ImmutableSet.of()
-        );
-
-        mService.agentDisconnected("com.android.foo");
-
-        verify(mOperationStorage, never())
-                .cancelOperation(anyInt(), anyBoolean(), any(IntConsumer.class));
-    }
-
-    @Test
     public void testReportDelayedRestoreResult_sendsLogsToMonitor() throws Exception {
         PackageInfo packageInfo = getPackageInfo(TEST_PACKAGE);
         when(mPackageManager.getPackageInfoAsUser(anyString(),
@@ -324,158 +284,6 @@
                 eq(packageInfo), eq(results), eq(OperationType.RESTORE));
     }
 
-    @Test
-    @DisableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_restrictedModeChangesFlagOff_shouldUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_keyValueBackup_shouldNotUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_keyValueRestore_shouldNotUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_RESTORE, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_packageOptedIn_shouldUseRestrictedMode()
-            throws Exception {
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenReturn(new PackageManager.Property(
-                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ true,
-                TEST_PACKAGE, /* className= */ null));
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_packageOptedOut_shouldNotUseRestrictedMode()
-            throws Exception {
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenReturn(new PackageManager.Property(
-                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ false,
-                TEST_PACKAGE, /* className= */ null));
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @DisableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_targetSdkBelowB_shouldUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_targetSdkB_notInList_shouldUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.clearNoRestrictedModePackages();
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_forRestore_targetSdkB_inList_shouldNotUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.RESTORE);
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_forBackup_targetSdkB_inList_shouldNotUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.BACKUP);
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
     private static PackageInfo getPackageInfo(String packageName) {
         PackageInfo packageInfo = new PackageInfo();
         packageInfo.applicationInfo = new ApplicationInfo();
@@ -487,8 +295,6 @@
         boolean isEnabledStatePersisted = false;
         boolean shouldUseNewBackupEligibilityRules = false;
 
-        private volatile Thread mWorkerThread = null;
-
         TestBackupService() {
             super(mContext, mPackageManager, mOperationStorage, mTransportManager, mBackupHandler,
                     createConstants(mContext), mActivityManager, mActivityManagerInternal);
@@ -523,26 +329,8 @@
         }
 
         @Override
-        Thread getThreadForAsyncOperation(String operationName, Runnable operation) {
-            mWorkerThread = super.getThreadForAsyncOperation(operationName, operation);
-            return mWorkerThread;
-        }
-
-        @Override
         BackupManagerMonitorEventSender getBMMEventSender(IBackupManagerMonitor monitor) {
             return mBackupManagerMonitorEventSender;
         }
-
-        private void waitForAsyncOperation() {
-            if (mWorkerThread == null) {
-                return;
-            }
-
-            try {
-                mWorkerThread.join(/* millis */ WORKER_THREAD_TIMEOUT_MILLISECONDS);
-            } catch (InterruptedException e) {
-                fail("Failed waiting for worker thread to complete: " + e.getMessage());
-            }
-        }
     }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
index 3310573..e618433 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
@@ -33,6 +33,7 @@
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.BackupAgentTimeoutParameters;
 import com.android.server.backup.OperationStorage;
 import com.android.server.backup.TransportManager;
@@ -66,6 +67,8 @@
     @Mock
     UserBackupManagerService mBackupManagerService;
     @Mock
+    BackupAgentConnectionManager mBackupAgentConnectionManager;
+    @Mock
     BackupTransportClient mBackupTransportClient;
     @Mock
     CountDownLatch mLatch;
@@ -95,6 +98,8 @@
         when(mBackupManagerService.isSetupComplete()).thenReturn(true);
         when(mBackupManagerService.getAgentTimeoutParameters()).thenReturn(
                 mBackupAgentTimeoutParameters);
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
         when(mBackupManagerService.getTransportManager()).thenReturn(mTransportManager);
         when(mTransportManager.getCurrentTransportClient(any())).thenReturn(mTransportConnection);
         when(mTransportConnection.connectOrThrow(any())).thenReturn(mBackupTransportClient);
@@ -142,11 +147,11 @@
 
         mTask.run();
 
-        InOrder inOrder = inOrder(mBackupManagerService);
-        inOrder.verify(mBackupManagerService).setNoRestrictedModePackages(
+        InOrder inOrder = inOrder(mBackupAgentConnectionManager);
+        inOrder.verify(mBackupAgentConnectionManager).setNoRestrictedModePackages(
                 eq(Set.of("package1")),
                 eq(BackupAnnotations.OperationType.BACKUP));
-        inOrder.verify(mBackupManagerService).clearNoRestrictedModePackages();
+        inOrder.verify(mBackupAgentConnectionManager).clearNoRestrictedModePackages();
     }
 
     private void createTask(String[] packageNames) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
index 055adf6..351aac3 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
@@ -44,6 +44,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.modules.utils.testing.TestableDeviceConfig;
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.Flags;
 import com.android.server.backup.UserBackupManagerService;
 import com.android.server.backup.internal.BackupHandler;
@@ -95,6 +96,8 @@
     private TransportConnection mTransportConnection;
     @Mock
     private BackupTransportClient mBackupTransportClient;
+    @Mock
+    private BackupAgentConnectionManager mBackupAgentConnectionManager;
 
     private Set<String> mExcludedkeys = new HashSet<>();
     private Map<String, String> mBackupData = new HashMap<>();
@@ -122,6 +125,9 @@
 
         mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
 
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
+
         mBackupDataSource = new ArrayDeque<>(mBackupData.keySet());
         when(mBackupDataInput.readNextHeader())
                 .then((Answer<Boolean>) invocation -> !mBackupDataSource.isEmpty());
@@ -166,7 +172,7 @@
         mRestoreTask.setNoRestrictedModePackages(mBackupTransportClient,
                 new PackageInfo[]{packageInfo1, packageInfo2});
 
-        verify(mBackupManagerService).setNoRestrictedModePackages(
+        verify(mBackupAgentConnectionManager).setNoRestrictedModePackages(
                 eq(Set.of("package1")),
                 eq(BackupAnnotations.OperationType.RESTORE));
     }