diff --git a/services/core/java/com/android/server/am/ActivityTaskManagerService.java b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
index 36261b5..b4f328b 100644
--- a/services/core/java/com/android/server/am/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
@@ -4846,8 +4846,7 @@
         updateResumedAppTrace(r);
         mLastResumedActivity = r;
 
-        // TODO(b/111361570): Support multiple focused apps in WM
-        mWindowManager.setFocusedApp(r.appToken, true);
+        r.getDisplay().setFocusedApp(r, true);
 
         applyUpdateLockStateLocked(r);
         applyUpdateVrModeLocked(r);