LLM-17026 [ai-completion] do not strand inline completion async log adds on awaiter cancellation

(cherry picked from commit 6a696ba70c3251af9188675eb772d13fcca2aa6e)

GitOrigin-RevId: d6ab7b100d7f0eb4e244f44944095d8e3663d279

platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt

@@ -72,8 +72,11 @@ class InlineCompletionLogsContainer() {
   private val asyncAdds = ConcurrentLinkedQueue<Job>()
 
   private suspend fun awaitAllAlreadyRunningAsyncAdds() {
-    while (currentCoroutineContext().isActive) {
-      val job = asyncAdds.poll() ?: return
+    // A snapshot on purpose: wait only for the adds that are already running, and never remove a job here.
+    // If this coroutine is canceled inside `join`, the job must stay visible to [cancelAsyncAdds], otherwise it keeps running
+    // on the application-level scope and retains everything its block captured (LLM-17026).
+    for (job in asyncAdds.toList()) {
+      if (!currentCoroutineContext().isActive) return
       job.join()
     }
   }
@@ -107,6 +110,9 @@ class InlineCompletionLogsContainer() {
       block().forEach { add(it) }
     }
     asyncAdds.add(job)
+    // [asyncAdds] is a registry of *pending* jobs: a finished job is useless for [cancelAsyncAdds], and only the producer may
+    // remove one, never an awaiter (see [awaitAllAlreadyRunningAsyncAdds]).
+    job.invokeOnCompletion { asyncAdds.remove(job) }
   }
 
   /**

platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt

@@ -10,6 +10,11 @@ import com.intellij.testFramework.ExtensionTestUtil
 import com.intellij.testFramework.LeakHunter
 import com.intellij.testFramework.LightPlatformTestCase
 import com.intellij.testFramework.common.timeoutRunBlocking
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.launch
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
@@ -194,6 +199,40 @@ class InlineCompletionLogsContainerTest : LightPlatformTestCase() {
     LeakHunter.checkLeak(logsContainer, Project::class.java)
   }
 
+  /**
+   * A canceled awaiter must not drop a still running async add from the cancellation registry:
+   * [InlineCompletionLogsContainer.logCurrent] must still be able to cancel it, otherwise the add keeps running on the
+   * application-level scope and retains everything its block captured. See LLM-17026.
+   */
+  @Test
+  fun testCanceledAwaitDoesNotStrandAsyncAdd(): Unit = timeoutRunBlocking {
+    val logsContainer = InlineCompletionLogsContainer()
+    logsContainer.mockRandom(1f)
+    val started = CompletableDeferred<Unit>()
+    val blockCanceled = CompletableDeferred<Unit>()
+    logsContainer.addAsync {
+      started.complete(Unit)
+      try {
+        awaitCancellation()
+      }
+      finally {
+        blockCanceled.complete(Unit)
+      }
+    }
+    started.await()
+
+    try {
+      // `Dispatchers.Unconfined` runs the awaiter eagerly on this thread up to its first real suspension,
+      // so it is guaranteed to be inside `join` by the time `launch` returns.
+      val awaiter = launch(Dispatchers.Unconfined) { logsContainer.awaitAndGetCurrentLogs() }
+      awaiter.cancelAndJoin()
+    }
+    finally {
+      logsContainer.logCurrent(project = null)
+    }
+    blockCanceled.await()
+  }
+
   private fun withEap(isEAP: Boolean, action: () -> Unit) {
     try {
       InlineCompletionEapSupport.getInstance().setMockEap(isEAP)