diff --git a/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl b/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
index a630873c7f677..37bd6d1c5843a 100644
--- a/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
+++ b/core/java/android/companion/ICompanionDeviceDiscoveryService.aidl
@@ -16,7 +16,6 @@
 
 package android.companion;
 
-import android.companion.Association;
 import android.companion.AssociationRequest;
 import android.companion.IFindDeviceCallback;
 import com.android.internal.infra.AndroidFuture;
@@ -28,5 +27,5 @@ oneway interface ICompanionDeviceDiscoveryService {
         in AssociationRequest request,
         in String callingPackage,
         in IFindDeviceCallback findCallback,
-        in AndroidFuture<Association> serviceCallback);
+        in AndroidFuture<String> serviceCallback);
 }
diff --git a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
index c24782e8b310f..3636f8f5c10e2 100644
--- a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
+++ b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionDeviceDiscoveryService.java
@@ -38,7 +38,6 @@ import android.bluetooth.le.ScanCallback;
 import android.bluetooth.le.ScanFilter;
 import android.bluetooth.le.ScanResult;
 import android.bluetooth.le.ScanSettings;
-import android.companion.Association;
 import android.companion.AssociationRequest;
 import android.companion.BluetoothDeviceFilter;
 import android.companion.BluetoothLeDeviceFilter;
@@ -95,7 +94,7 @@ public class CompanionDeviceDiscoveryService extends Service {
     DeviceFilterPair mSelectedDevice;
     IFindDeviceCallback mFindCallback;
 
-    AndroidFuture<Association> mServiceCallback;
+    AndroidFuture<String> mServiceCallback;
     boolean mIsScanning = false;
     @Nullable
     CompanionDeviceActivity mActivity = null;
@@ -106,7 +105,7 @@ public class CompanionDeviceDiscoveryService extends Service {
         public void startDiscovery(AssociationRequest request,
                 String callingPackage,
                 IFindDeviceCallback findCallback,
-                AndroidFuture serviceCallback) {
+                AndroidFuture<String> serviceCallback) {
             Log.i(LOG_TAG,
                     "startDiscovery() called with: filter = [" + request
                             + "], findCallback = [" + findCallback + "]"
@@ -320,9 +319,7 @@ public class CompanionDeviceDiscoveryService extends Service {
         if (callingPackage == null || deviceAddress == null) {
             return;
         }
-        mServiceCallback.complete(new Association(
-                getUserId(), deviceAddress, callingPackage, mRequest.getDeviceProfile(), false,
-                System.currentTimeMillis()));
+        mServiceCallback.complete(deviceAddress);
     }
 
     void onCancel() {
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 44a49978795c7..34d9c8adc3249 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -208,7 +208,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
     private ScanCallback mBleScanCallback = new BleScanCallback();
     private AssociationRequest mRequest;
     private String mCallingPackage;
-    private AndroidFuture<Association> mOngoingDeviceDiscovery;
+    private AndroidFuture<?> mOngoingDeviceDiscovery;
     private PermissionControllerManager mPermissionControllerManager;
 
     private BluetoothDeviceConnectedListener mBluetoothDeviceConnectedListener =
@@ -383,7 +383,7 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
         Slog.d(LOG_TAG, "cleanup(); discovery = "
                 + mOngoingDeviceDiscovery + ", request = " + mRequest);
         synchronized (mLock) {
-            AndroidFuture<Association> ongoingDeviceDiscovery = mOngoingDeviceDiscovery;
+            AndroidFuture<?> ongoingDeviceDiscovery = mOngoingDeviceDiscovery;
             if (ongoingDeviceDiscovery != null && !ongoingDeviceDiscovery.isDone()) {
                 ongoingDeviceDiscovery.cancel(true);
             }
@@ -458,13 +458,16 @@ public class CompanionDeviceManagerService extends SystemService implements Bind
                 return mServiceConnectors.forUser(userId).postAsync(service -> {
                     Slog.d(LOG_TAG, "Connected to CDM service; starting discovery for " + request);
 
-                    AndroidFuture<Association> future = new AndroidFuture<>();
+                    AndroidFuture<String> future = new AndroidFuture<>();
                     service.startDiscovery(request, callingPackage, callback, future);
                     return future;
                 }).cancelTimeout();
 
-            }, FgThread.getExecutor()).whenComplete(uncheckExceptions((association, err) -> {
+            }, FgThread.getExecutor()).whenComplete(uncheckExceptions((deviceAddress, err) -> {
                 if (err == null) {
+                    Association association = new Association(userId, deviceAddress, callingPackage,
+                            mRequest.getDeviceProfile(), false,
+                            System.currentTimeMillis());
                     addAssociation(association, userId);
                 } else {
                     Slog.e(LOG_TAG, "Failed to discover device(s)", err);