From 1a1f6849ac4af6d7609873560e160a101b9fdbc6 Mon Sep 17 00:00:00 2001
From: "Kirill.Karnaukhov" <kirill.karnaukhov@jetbrains.com>
Date: Mon, 28 Oct 2024 15:23:21 +0100
Subject: [PATCH] [inline-completion] IJPL-165258: do not destroy Handler state
 when something unexpected happens to Session

GitOrigin-RevId: 5d2b5aed72da4710271f8bb8ce589e0238c07449
---
 .../completion/InlineCompletionHandler.kt     | 11 ++++-
 .../session/InlineCompletionSession.kt        |  9 +++-
 .../InlineCompletionExceptionTest.kt          | 45 +++++++++++++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/InlineCompletionExceptionTest.kt

diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionHandler.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionHandler.kt
index b12ca7c2896f..9c56680d0a8c 100644
--- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionHandler.kt
+++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/InlineCompletionHandler.kt
@@ -218,7 +218,16 @@ abstract class InlineCompletionHandler @ApiStatus.Internal constructor(
   @RequiresEdt
   fun hide(context: InlineCompletionContext, finishType: FinishType = FinishType.OTHER) {
     ThreadingAssertions.assertEventDispatchThread()
-    LOG.assertTrue(!context.isDisposed)
+    val currentSession = InlineCompletionSession.getOrNull(editor)
+    if (context !== currentSession?.context) {
+      LOG.error("[Inline Completion] Cannot hide a session because an invalid context is passed.")
+      return
+    }
+    if (context.isDisposed) {
+      sessionManager.removeSession()
+      LOG.error("[Inline Completion] Cannot hide a session because the context is already disposed.")
+      return
+    }
     doHide(context, finishType)
   }
 
diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/session/InlineCompletionSession.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/session/InlineCompletionSession.kt
index 8922939a136c..f41f16fb5382 100644
--- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/session/InlineCompletionSession.kt
+++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/session/InlineCompletionSession.kt
@@ -192,9 +192,14 @@ class InlineCompletionSession private constructor(
     @RequiresEdt
     internal fun remove(editor: Editor) {
       getOrNull(editor)?.apply {
-        Disposer.dispose(this)
+        if (!context.isDisposed) {
+          Disposer.dispose(this)
+          LOG.trace("[Inline Completion] Remove inline completion session")
+        }
+        else {
+          LOG.warn("[Inline Completion] Cannot dispose session because it's already disposed.")
+        }
         editor.putUserData(INLINE_COMPLETION_SESSION, null)
-        LOG.trace("Remove inline completion session")
       }
     }
   }
diff --git a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/InlineCompletionExceptionTest.kt b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/InlineCompletionExceptionTest.kt
new file mode 100644
index 000000000000..55a6410475d4
--- /dev/null
+++ b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/InlineCompletionExceptionTest.kt
@@ -0,0 +1,45 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.inline.completion
+
+import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement
+import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement
+import com.intellij.codeInsight.inline.completion.impl.SimpleInlineCompletionProvider
+import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession
+import com.intellij.openapi.application.EDT
+import com.intellij.openapi.fileTypes.PlainTextFileType
+import com.intellij.openapi.progress.assertLogThrows
+import com.intellij.openapi.util.Disposer
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+internal class InlineCompletionExceptionTest : InlineCompletionTestCase() {
+
+  private fun registerSuggestion(vararg suggestion: InlineCompletionElement) {
+    InlineCompletionHandler.registerTestHandler(SimpleInlineCompletionProvider(suggestion.toList()), testRootDisposable)
+  }
+
+  @Test
+  fun `test manually dispose session and use completion after`() = myFixture.testInlineCompletion {
+    init(PlainTextFileType.INSTANCE)
+    registerSuggestion(InlineCompletionGrayTextElement("simple completion"))
+    typeChar('\n')
+    delay()
+    withContext(Dispatchers.EDT) {
+      val session = assertNotNull(InlineCompletionSession.getOrNull(myFixture.editor))
+      Disposer.dispose(session)
+    }
+    withContext(Dispatchers.EDT) {
+      assertLogThrows<Throwable> {
+        escape()
+      }
+    }
+    assertInlineHidden()
+    typeChar('\n')
+    delay()
+    assertInlineRender("simple completion")
+  }
+}