From 07f06c698694ad13b74aeff365804b97c201099c Mon Sep 17 00:00:00 2001
From: Kirill Karnaukhov <kirill.karnaukhov@jetbrains.com>
Date: Wed, 12 Aug 2026 16:28:15 +0200
Subject: [PATCH] LLM-17026 [ai-completion] prevent async log adds after flush

(cherry picked from commit 1f8f3ad40eb0e92883f4c6332a6dfb419807a663)

GitOrigin-RevId: 0f7971dc7598689923ad7cf8c68de831b9dcfd35
---
 .../logs/InlineCompletionLogsContainer.kt     | 16 +++++++++++++-
 .../logs/InlineCompletionLogsContainerTest.kt | 22 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt b/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt
index d4603f05aef1..f3049fdbe9c5 100644
--- a/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt
+++ b/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt
@@ -12,6 +12,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.removeUserData
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.isActive
@@ -70,6 +71,7 @@ class InlineCompletionLogsContainer() {
   }
 
   private val asyncAdds = ConcurrentLinkedQueue<Job>()
+  private val asyncAddsClosed = AtomicBoolean(false)
 
   private suspend fun awaitAllAlreadyRunningAsyncAdds() {
     // A snapshot on purpose: wait only for the adds that are already running, and never remove a job here.
@@ -82,6 +84,8 @@ class InlineCompletionLogsContainer() {
   }
 
   private fun cancelAsyncAdds() {
+    // Close first, so a concurrent addAsync either stays visible in the queue or cancels itself before it can start.
+    asyncAddsClosed.set(true)
     while (true) {
       val job = asyncAdds.poll() ?: break
       job.cancel()
@@ -104,15 +108,25 @@ class InlineCompletionLogsContainer() {
 
   /**
    * Use [add] if there is no special need to use async variant. See [add] documentation for more info.
+   * Calls made after [logCurrent] are ignored because the container has already been finalized.
    */
   fun addAsync(block: suspend () -> List<EventPair<*>>) {
-    val job = InlineCompletionLogsScopeProvider.getInstance().cs.launch {
+    if (asyncAddsClosed.get()) return
+
+    // Start lazily so every running job is already visible to cancelAsyncAdds.
+    val job = InlineCompletionLogsScopeProvider.getInstance().cs.launch(start = CoroutineStart.LAZY) {
       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) }
+    if (asyncAddsClosed.get()) {
+      job.cancel()
+    }
+    else {
+      job.start()
+    }
   }
 
   /**
diff --git a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt
index 715b4d57545b..8e5134a8973a 100644
--- a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt
+++ b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt
@@ -233,6 +233,26 @@ class InlineCompletionLogsContainerTest : LightPlatformTestCase() {
     blockCanceled.await()
   }
 
+  /**
+   * Once the container has been finalized, a late async add must not start: there is no later [InlineCompletionLogsContainer.logCurrent]
+   * call that could cancel it or send its result.
+   */
+  @Test
+  fun testAsyncAddAfterLogCurrentDoesNotStart(): Unit = timeoutRunBlocking {
+    val logsContainer = InlineCompletionLogsContainer()
+    logsContainer.mockRandom(1f)
+    logsContainer.logCurrent(project = null)
+    val started = CompletableDeferred<Unit>()
+
+    logsContainer.addAsync {
+      started.complete(Unit)
+      listOf(TestPhasedLogs.basicTestField with 42)
+    }
+    logsContainer.awaitAndGetCurrentLogs()
+
+    assertFalse("An async add submitted after logCurrent must not start", started.isCompleted)
+  }
+
   private fun withEap(isEAP: Boolean, action: () -> Unit) {
     try {
       InlineCompletionEapSupport.getInstance().setMockEap(isEAP)
@@ -276,4 +296,4 @@ object TestPhasedLogs : PhasedLogs(Phase.INLINE_API_STARTING) {
 
 class TestInlineCompletionSessionLogsEP : InlineCompletionSessionLogsEP {
   override val logGroups: List<PhasedLogs> = listOf(TestPhasedLogs)
-}
\ No newline at end of file
+}