Change certain guarantees regarding TextInputSession

`InputConnection#closeConnection` and `View#onCreateInputConnection` no longer close an existing `PlatformTextInputSession`. Input Method Framework does not have any guarantees in terms of when these functions are called and their exact meaning. Only guarantee is that an `InputConnection` can be removed if it receives a `closeConnection` call. It doesn't necessarily mean anything about the ongoing input session.

Test: gradle :compose:ui:ui:cAT
Test: gradle :compose:foundation:foundation:cAT
Bug: 333867282
Fix: 349117269
(cherry picked from https://android-review.googlesource.com/q/commit:d15e3fa93c25c8b889ba12d214fe29797002bcb3)
(cherry picked from https://android-review.googlesource.com/q/commit:32d3c166432ca139b0759af7494cd9a861efefdf)
Merged-In: Ic017e8dfea484fa9461af915c19f6481a654bd0c
Change-Id: Ic017e8dfea484fa9461af915c19f6481a654bd0c
diff --git a/compose/ui/ui/src/androidInstrumentedTest/kotlin/androidx/compose/ui/text/input/PlatformTextInputViewIntegrationTest.kt b/compose/ui/ui/src/androidInstrumentedTest/kotlin/androidx/compose/ui/text/input/PlatformTextInputViewIntegrationTest.kt
index 8eba172..ca0f1f6 100644
--- a/compose/ui/ui/src/androidInstrumentedTest/kotlin/androidx/compose/ui/text/input/PlatformTextInputViewIntegrationTest.kt
+++ b/compose/ui/ui/src/androidInstrumentedTest/kotlin/androidx/compose/ui/text/input/PlatformTextInputViewIntegrationTest.kt
@@ -472,7 +472,7 @@
     // closeConnection is only supported on API 24+
     @SdkSuppress(minSdkVersion = 24)
     @Test
-    fun connectionClosed_whenCreateConnectionCalledAgain() {
+    fun connectionNotClosed_whenCreateConnectionCalledAgain() {
         class TestConnection(view: View) : BaseInputConnection(view, true) {
             var closeCalls = 0
             override fun closeConnection() {
@@ -504,15 +504,15 @@
 
             assertThat(connections).hasSize(2)
             val connection2 = connections.last()
-            assertThat(connection1.closeCalls).isEqualTo(1)
+            assertThat(connection1.closeCalls).isEqualTo(0)
             assertThat(connection2.closeCalls).isEqualTo(0)
 
             hostView.onCreateInputConnection(EditorInfo())
 
             assertThat(connections).hasSize(3)
             val connection3 = connections.last()
-            assertThat(connection1.closeCalls).isEqualTo(1)
-            assertThat(connection2.closeCalls).isEqualTo(1)
+            assertThat(connection1.closeCalls).isEqualTo(0)
+            assertThat(connection2.closeCalls).isEqualTo(0)
             assertThat(connection3.closeCalls).isEqualTo(0)
         }
 
@@ -523,7 +523,7 @@
 
     @SdkSuppress(minSdkVersion = 24)
     @Test
-    fun innerSessionCanceled_whenIsolatedFromOuterSession_whenConnectionClosed() {
+    fun innerSessionNotCanceled_whenIsolatedFromOuterSession_whenConnectionClosed() {
         setupContent()
         lateinit var innerJob: Job
         val outerJob = coroutineScope.launch {
@@ -543,15 +543,14 @@
         }
 
         rule.runOnIdle {
-            // Outer job isn't canceled, only the inner one.
             assertThat(outerJob.isActive).isTrue()
-            assertThat(innerJob.isActive).isFalse()
+            assertThat(innerJob.isActive).isTrue()
         }
     }
 
     @SdkSuppress(minSdkVersion = 24)
     @Test
-    fun cancellationPropagates_whenConnectionClosed() {
+    fun cancellationDoesNotPropagate_whenConnectionClosed() {
         setupContent()
         val sessionJob = coroutineScope.launch {
             node1.establishTextInputSession {
@@ -565,9 +564,7 @@
             connection.closeConnection()
         }
 
-        rule.runOnIdle {
-            assertThat(sessionJob.isActive).isFalse()
-        }
+        rule.runOnIdle { assertThat(sessionJob.isActive).isTrue() }
     }
 
     @Test
diff --git a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidPlatformTextInputSession.android.kt b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidPlatformTextInputSession.android.kt
index ed817a3..b5fa2dc 100644
--- a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidPlatformTextInputSession.android.kt
+++ b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/AndroidPlatformTextInputSession.android.kt
@@ -21,9 +21,11 @@
 import android.view.View
 import android.view.inputmethod.EditorInfo
 import android.view.inputmethod.InputConnection
+import androidx.compose.runtime.collection.mutableVectorOf
 import androidx.compose.ui.InternalComposeUiApi
 import androidx.compose.ui.SessionMutex
 import androidx.compose.ui.node.Owner
+import androidx.compose.ui.node.WeakReference
 import androidx.compose.ui.text.InternalTextApi
 import androidx.compose.ui.text.input.NullableInputConnectionWrapper
 import androidx.compose.ui.text.input.TextInputService
@@ -69,10 +71,11 @@
 
     override suspend fun startInputMethod(request: PlatformTextInputMethodRequest): Nothing =
         methodSessionMutex.withSessionCancellingPrevious(
-            sessionInitializer = { coroutineScope ->
-                InputMethodSession(request, onConnectionClosed = {
-                    coroutineScope.cancel()
-                })
+            sessionInitializer = {
+                InputMethodSession(
+                    request = request,
+                    onAllConnectionsClosed = { coroutineScope.cancel() }
+                )
             }
         ) { methodSession ->
             @Suppress("RemoveExplicitTypeArguments")
@@ -109,13 +112,16 @@
  * [AndroidPlatformTextInputSession]'s input method session. Instances of this class correspond to
  * calls to [AndroidPlatformTextInputSession.startInputMethod]. This class ensures that old
  * connections are disposed before new ones are created.
+ *
+ * @param onAllConnectionsClosed Called when all created [InputConnection]s receive
+ *   [InputConnection.closeConnection] call.
  */
 private class InputMethodSession(
     private val request: PlatformTextInputMethodRequest,
-    private val onConnectionClosed: () -> Unit
+    private val onAllConnectionsClosed: () -> Unit
 ) {
     private val lock = Any()
-    private var connection: NullableInputConnectionWrapper? = null
+    private var connections = mutableVectorOf<WeakReference<NullableInputConnectionWrapper>>()
     private var disposed = false
 
     val isActive: Boolean get() = !disposed
@@ -130,29 +136,48 @@
     fun createInputConnection(outAttrs: EditorInfo): InputConnection? {
         synchronized(lock) {
             if (disposed) return null
-            // Manually close the delegate in case the system won't until later.
-            connection?.disposeDelegate()
+
+            // Do not manually dispose a previous InputConnection until it's collected by the GC or
+            // an explicit call is received to its `onConnectionClosed` callback.
 
             val connectionDelegate = request.createInputConnection(outAttrs)
             return NullableInputConnectionWrapper(
-                delegate = connectionDelegate,
-                onConnectionClosed = onConnectionClosed
-            ).also {
-                connection = it
-            }
+                    delegate = connectionDelegate,
+                    onConnectionClosed = { closedConnection ->
+                        // We should not cancel any ongoing input session because connection is
+                        // closed
+                        // from InputConnection. This may happen at any time and does not indicate
+                        // whether the system is trying to stop an input session. The platform may
+                        // request a new InputConnection immediately after closing this one.
+                        // Instead we should just clear all the resources used by this
+                        // inputConnection,
+                        // because the platform guarantees that it will never reuse a closed
+                        // connection.
+                        closedConnection.disposeDelegate()
+                        val removeIndex = connections.indexOfFirst { it == closedConnection }
+                        if (removeIndex >= 0) connections.removeAt(removeIndex)
+                        if (connections.isEmpty()) {
+                            onAllConnectionsClosed()
+                        }
+                    }
+                )
+                .also { connections.add(WeakReference(it)) }
         }
     }
 
     /**
      * Disposes the current [InputConnection]. After calling this method, all future calls to
      * [createInputConnection] will return null.
+     *
+     * This function is only called from coroutine cancellation routine so it's not required to
+     * cancel the coroutine from here.
      */
     fun dispose() {
         synchronized(lock) {
             // Manually close the delegate in case the system forgets to.
             disposed = true
-            connection?.disposeDelegate()
-            connection = null
+            connections.forEach { it.get()?.disposeDelegate() }
+            connections.clear()
         }
     }
 }
diff --git a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/PlatformTextInputMethodRequest.android.kt b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/PlatformTextInputMethodRequest.android.kt
index 2813d78..5bafc46 100644
--- a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/PlatformTextInputMethodRequest.android.kt
+++ b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/platform/PlatformTextInputMethodRequest.android.kt
@@ -29,26 +29,30 @@
     /**
      * Called when the platform requests an [InputConnection] via [View.onCreateInputConnection].
      *
-     * This method makes stricter ordering guarantees about the lifetime of the returned
+     * This method makes relatively stricter ordering guarantees about the lifetime of the returned
      * [InputConnection] than Android does, to make working with connections simpler. Namely, it
      * guarantees:
-     *  - For a given [PlatformTextInputMethodRequest] instance, only one [InputConnection] will
-     *    ever be active at a time.
-     *  - References to an [InputConnection] will be cleared as soon as the connection becomes
-     *    inactive. Even if Android leaks its reference to the connection, the connection returned
-     *    from this method will not be leaked.
-     *  - On API levels that support [InputConnection.closeConnection] (24+), a connection will
-     *    always be closed before a new connection is requested.
+     * - References to an [InputConnection] will be cleared as soon as the connection get closed via
+     *   [InputConnection.closeConnection]. Even if Android leaks its reference to the connection,
+     *   the connection returned from this method will not be leaked.
+     *
+     * However it does not guarantee that only one [InputConnection] will ever be active at a time
+     * for a given [PlatformTextInputMethodRequest] instance. On the other hand Android platform
+     * does guarantee that even though the platform may create multiple [InputConnection]s, only one
+     * of them will ever communicate with the app, invalidating any other [InputConnection] that
+     * remained open at the time of communication.
      *
      * Android may call [View.onCreateInputConnection] multiple times for the same session – each
-     * system call will result in a 1:1 call to this method, although the old connection will always
-     * be closed first.
+     * system call will result in a 1:1 call to this method. Unfortunately Android platform may
+     * decide to use an earlier [InputConnection] returned from this function, invalidating the ones
+     * that were created later. Please do not rely on the order of calls to this function.
      *
      * @param outAttributes The [EditorInfo] from [View.onCreateInputConnection].
-     *
      * @return The [InputConnection] that will be used to talk to the IME as long as the session is
-     * active. This connection will not receive any calls after the requesting coroutine is
-     * cancelled.
+     *   active. This connection will not receive any calls after the requesting coroutine is
+     *   cancelled.
      */
+    // Please take a look at go/text-input-session-android-gotchas to learn more about corner cases
+    // of Android InputConnection management
     fun createInputConnection(outAttributes: EditorInfo): InputConnection
 }
diff --git a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/text/input/NullableInputConnectionWrapper.android.kt b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/text/input/NullableInputConnectionWrapper.android.kt
index 522b0171..75fe3ac 100644
--- a/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/text/input/NullableInputConnectionWrapper.android.kt
+++ b/compose/ui/ui/src/androidMain/kotlin/androidx/compose/ui/text/input/NullableInputConnectionWrapper.android.kt
@@ -39,33 +39,23 @@
  * @param delegate The [InputConnection] that will receive all calls on this object until
  * `disposeDelegate` or `closeConnection` are called.
  * @param onConnectionClosed A callback that will be invoked the first time `closeConnection` is
- * called. Will not be invoked by `disposeDelegate`, and will not be invoked if `disposeDelegate` is
- * called before `closeConnection`.
+ *   called. Will not be invoked by `disposeDelegate`, and will not be invoked if `disposeDelegate`
+ *   is called before `closeConnection`. Passed in [NullableInputConnectionWrapper] is the same
+ *   instance that is created here.
  */
 internal fun NullableInputConnectionWrapper(
     delegate: InputConnection,
-    onConnectionClosed: () -> Unit
-): NullableInputConnectionWrapper = when {
-    Build.VERSION.SDK_INT >= 34 -> NullableInputConnectionWrapperApi34(
-        delegate,
-        onConnectionClosed
-    )
-
-    Build.VERSION.SDK_INT >= 25 -> NullableInputConnectionWrapperApi25(
-        delegate,
-        onConnectionClosed
-    )
-
-    Build.VERSION.SDK_INT >= 24 -> NullableInputConnectionWrapperApi24(
-        delegate,
-        onConnectionClosed
-    )
-
-    else -> NullableInputConnectionWrapperApi21(
-        delegate,
-        onConnectionClosed
-    )
-}
+    onConnectionClosed: (NullableInputConnectionWrapper) -> Unit
+): NullableInputConnectionWrapper =
+    when {
+        Build.VERSION.SDK_INT >= 34 ->
+            NullableInputConnectionWrapperApi34(delegate, onConnectionClosed)
+        Build.VERSION.SDK_INT >= 25 ->
+            NullableInputConnectionWrapperApi25(delegate, onConnectionClosed)
+        Build.VERSION.SDK_INT >= 24 ->
+            NullableInputConnectionWrapperApi24(delegate, onConnectionClosed)
+        else -> NullableInputConnectionWrapperApi21(delegate, onConnectionClosed)
+    }
 
 /**
  * An [InputConnection] that will delegate all calls to a delegate [InputConnection]. This is
@@ -93,7 +83,7 @@
 
 private open class NullableInputConnectionWrapperApi21(
     delegate: InputConnection,
-    private val onConnectionClosed: () -> Unit
+    private val onConnectionClosed: (NullableInputConnectionWrapper) -> Unit
 ) : NullableInputConnectionWrapper {
 
     protected var delegate: InputConnection? = delegate
@@ -109,7 +99,7 @@
     final override fun closeConnection() {
         delegate?.let {
             disposeDelegate()
-            onConnectionClosed()
+            onConnectionClosed(this)
         }
     }
 
@@ -189,7 +179,7 @@
 @RequiresApi(24)
 private open class NullableInputConnectionWrapperApi24(
     delegate: InputConnection,
-    onConnectionClosed: () -> Unit
+    onConnectionClosed: (NullableInputConnectionWrapper) -> Unit
 ) : NullableInputConnectionWrapperApi21(delegate, onConnectionClosed) {
 
     final override fun deleteSurroundingTextInCodePoints(p0: Int, p1: Int): Boolean =
@@ -205,7 +195,7 @@
 @RequiresApi(25)
 private open class NullableInputConnectionWrapperApi25(
     delegate: InputConnection,
-    onConnectionClosed: () -> Unit
+    onConnectionClosed: (NullableInputConnectionWrapper) -> Unit
 ) : NullableInputConnectionWrapperApi24(delegate, onConnectionClosed) {
 
     final override fun commitContent(p0: InputContentInfo, p1: Int, p2: Bundle?): Boolean =
@@ -215,7 +205,7 @@
 @RequiresApi(34)
 private open class NullableInputConnectionWrapperApi34(
     delegate: InputConnection,
-    onConnectionClosed: () -> Unit
+    onConnectionClosed: (NullableInputConnectionWrapper) -> Unit
 ) : NullableInputConnectionWrapperApi25(delegate, onConnectionClosed) {
     final override fun performHandwritingGesture(
         gesture: HandwritingGesture,
diff --git a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/platform/PlatformTextInputModifierNode.kt b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/platform/PlatformTextInputModifierNode.kt
index 0b34c12..3a6917f 100644
--- a/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/platform/PlatformTextInputModifierNode.kt
+++ b/compose/ui/ui/src/commonMain/kotlin/androidx/compose/ui/platform/PlatformTextInputModifierNode.kt
@@ -113,11 +113,15 @@
  * cancellation tasks (e.g. `finally` blocks) before running the new [block] function.
  *
  * The session will be closed when:
- *  - The session function throws an exception.
- *  - The requesting coroutine is cancelled.
- *  - Another session is started via this method, either from the same modifier or a different one.
- *  - The system closes the connection (currently only supported on Android, and there only
- *    depending on OS version).
+ * - The session function throws an exception.
+ * - The requesting coroutine is cancelled.
+ * - Another session is started via this method, either from the same modifier or a different one.
+ *
+ * The session may remain open when:
+ * - The system closes the connection. This behavior currently only exists on Android depending on
+ *   OS version. Android platform may intermittently close the active connection to immediately
+ *   start it back again. In these cases the session will not be prematurely closed, so that it can
+ *   serve the follow-up requests.
  *
  * This function should only be called from the modifier node's
  * [coroutineScope][Modifier.Node.coroutineScope]. If it is not, the session will _not_