diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index faa98fa..9a8824f 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -443,6 +443,9 @@
     // VoiceInteractionManagerService
     ComponentName mActiveVoiceInteractionServiceComponent;
 
+    // A map userId and all its companion app packages
+    private final Map<Integer, Set<String>> mCompanionAppPackageMap = new ArrayMap<>();
+
     VrController mVrController;
     KeyguardController mKeyguardController;
     private final ClientLifecycleManager mLifecycleManager;
@@ -5909,6 +5912,14 @@
         }
     }
 
+    boolean isAssociatedCompanionApp(int userId, String packageName) {
+        final Set<String> allPackages = mCompanionAppPackageMap.get(userId);
+        if (allPackages == null) {
+            return false;
+        }
+        return allPackages.contains(packageName);
+    }
+
     final class H extends Handler {
         static final int REPORT_TIME_TRACKER_MSG = 1;
 
@@ -7282,5 +7293,17 @@
                 ActivityTaskManagerService.this.setDeviceOwnerPackageName(deviceOwnerPkg);
             }
         }
+
+        @Override
+        public void setCompanionAppPackages(int userId, Set<String> companionAppPackages) {
+            // Deep copy all content to make sure we do not rely on the source
+            final Set<String> result = new HashSet<>();
+            for (String pkg : companionAppPackages) {
+                result.add(pkg);
+            }
+            synchronized (mGlobalLock) {
+                mCompanionAppPackageMap.put(userId, result);
+            }
+        }
     }
 }