Correctly setup IME subtypes in KeyboardLayoutManager in multi-user case

Bug: 354333072
Test: atest KeyboardLayoutManagerTests
Flag: com.android.hardware.input.keyboard_layout_manager_multi_user_ime_setup
Change-Id: I494a2ab5bd7c2fead341a9821144594efb068909

2 files changed, 21 insertions, 3 deletions

diff --git a/core/java/android/hardware/input/input_framework.aconfig b/core/java/android/hardware/input/input_framework.aconfig
index 0cd280002dbf..b4ad050c8b5d 100644
--- a/core/java/android/hardware/input/input_framework.aconfig
+++ b/core/java/android/hardware/input/input_framework.aconfig
@@ -75,3 +75,13 @@ flag {
     description: "Enables a developer overlay that displays raw touchpad input data and gesture recognition status in real-time."
     bug: "286551975"
 }
+
+flag {
+    namespace: "input_native"
+    name: "keyboard_layout_manager_multi_user_ime_setup"
+    description: "Update KeyboardLayoutManager to work correctly with multi-user IME setup"
+    bug: "354333072"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/services/core/java/com/android/server/input/KeyboardLayoutManager.java b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
index 49934126ab8c..1d1a178ff20b 100644
--- a/services/core/java/com/android/server/input/KeyboardLayoutManager.java
+++ b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
@@ -22,6 +22,8 @@ import static android.hardware.input.KeyboardLayoutSelectionResult.LAYOUT_SELECT
 import static android.hardware.input.KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_VIRTUAL_KEYBOARD;
 import static android.hardware.input.KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_DEFAULT;
 
+import static com.android.hardware.input.Flags.keyboardLayoutManagerMultiUserImeSetup;
+
 import android.annotation.AnyThread;
 import android.annotation.MainThread;
 import android.annotation.NonNull;
@@ -1066,9 +1068,15 @@ class KeyboardLayoutManager implements InputManager.InputDeviceListener {
             for (InputMethodInfo imeInfo :
                     inputMethodManagerInternal.getEnabledInputMethodListAsUser(
                             userId)) {
-                for (InputMethodSubtype imeSubtype :
-                        inputMethodManager.getEnabledInputMethodSubtypeList(
-                                imeInfo, true /* allowsImplicitlyEnabledSubtypes */)) {
+                final List<InputMethodSubtype> imeSubtypes;
+                if (keyboardLayoutManagerMultiUserImeSetup()) {
+                    imeSubtypes = inputMethodManagerInternal.getEnabledInputMethodSubtypeListAsUser(
+                            imeInfo.getId(), true /* allowsImplicitlyEnabledSubtypes */, userId);
+                } else {
+                    imeSubtypes = inputMethodManager.getEnabledInputMethodSubtypeList(imeInfo,
+                            true /* allowsImplicitlyEnabledSubtypes */);
+                }
+                for (InputMethodSubtype imeSubtype : imeSubtypes) {
                     if (!imeSubtype.isSuitableForPhysicalKeyboardLayoutMapping()) {
                         continue;
                     }