From c16b5740556d39e9cb2323ffdc164475bd757883 Mon Sep 17 00:00:00 2001
From: Shan Huang <shanh@google.com>
Date: Tue, 10 May 2022 16:09:13 +0000
Subject: [PATCH] Fix IME callback not being re-registered on screen off / on.

Bug:231917948
Bug:229355440
Test: atest InputMethodServiceLifecycleTest
Test: atest CtsInputMethodTestCases:InputMethodServiceTest
Test atest CtsInputMethodTestCases:KeyboardVisibilityControlTest

Change-Id: I8dc1425aa9338b37128d308bf58fd15c39d04a88
---
 .../java/android/inputmethodservice/InputMethodService.java | 6 +++++-
 core/java/android/window/ImeOnBackInvokedDispatcher.java    | 6 ++++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 25296bc0a8b91..200fe22edaadd 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -1625,7 +1625,6 @@ public class InputMethodService extends AbstractInputMethodService {
             // when IME developers are doing something unsupported.
             InputMethodPrivilegedOperationsRegistry.remove(mToken);
         }
-        unregisterCompatOnBackInvokedCallback();
         mImeDispatcher = null;
     }
 
@@ -2788,6 +2787,11 @@ public class InputMethodService extends AbstractInputMethodService {
         if (mInkWindow != null) {
             finishStylusHandwriting();
         }
+        // Back callback is typically unregistered in {@link #hideWindow()}, but it's possible
+        // for {@link #doFinishInput()} to be called without {@link #hideWindow()} so we also
+        // unregister here.
+        // TODO(b/232341407): Add CTS to verify back behavior after screen on / off.
+        unregisterCompatOnBackInvokedCallback();
     }
 
     void doStartInput(InputConnection ic, EditorInfo attribute, boolean restarting) {
diff --git a/core/java/android/window/ImeOnBackInvokedDispatcher.java b/core/java/android/window/ImeOnBackInvokedDispatcher.java
index d5763aa25884f..8bdf233b81c8a 100644
--- a/core/java/android/window/ImeOnBackInvokedDispatcher.java
+++ b/core/java/android/window/ImeOnBackInvokedDispatcher.java
@@ -158,6 +158,12 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
 
     /** Clears all registered callbacks on the instance. */
     public void clear() {
+        // Unregister previously registered callbacks if there's any.
+        if (getReceivingDispatcher() != null) {
+            for (OnBackInvokedCallback callback : mImeCallbackMap.values()) {
+                getReceivingDispatcher().unregisterOnBackInvokedCallback(callback);
+            }
+        }
         mImeCallbackMap.clear();
     }