fix(amazonq): fix UI freeze when editing large files (#5685)

rli · web-flow · commit 30f3ccac772d · 2025-05-09T13:08:02.000-07:00
diff --git a/.changes/next-release/bugfix-1d632e06-24fa-40f1-a8ae-33cca0d86ec4.json b/.changes/next-release/bugfix-1d632e06-24fa-40f1-a8ae-33cca0d86ec4.json
@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix UI freezes that may occur when interacting with large files in the editor"
+}
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandler.kt
@@ -4,7 +4,10 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.textdocument
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.event.DocumentEvent
+import com.intellij.openapi.editor.event.DocumentListener
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileDocumentManagerListener
 import com.intellij.openapi.fileEditor.FileEditorManager
@@ -29,10 +32,11 @@ import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 
 class TextDocumentServiceHandler(
     private val project: Project,
-    serverInstance: Disposable,
+    private val serverInstance: Disposable,
 ) : FileDocumentManagerListener,
     FileEditorManagerListener,
-    BulkFileListener {
+    BulkFileListener,
+    DocumentListener {
 
     init {
         // didOpen & didClose events
@@ -61,18 +65,30 @@ class TextDocumentServiceHandler(
     }
 
     private fun handleFileOpened(file: VirtualFile) {
+        ApplicationManager.getApplication().runReadAction {
+            FileDocumentManager.getInstance().getDocument(file)?.addDocumentListener(
+                object : DocumentListener {
+                    override fun documentChanged(event: DocumentEvent) {
+                        realTimeEdit(event)
+                    }
+                },
+                serverInstance
+            )
+        }
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             toUriString(file)?.let { uri ->
-                languageServer.textDocumentService.didOpen(
-                    DidOpenTextDocumentParams().apply {
-                        textDocument = TextDocumentItem().apply {
-                            this.uri = uri
-                            text = file.inputStream.readAllBytes().decodeToString()
-                            languageId = file.fileType.name.lowercase()
-                            version = file.modificationStamp.toInt()
+                pluginAwareExecuteOnPooledThread {
+                    languageServer.textDocumentService.didOpen(
+                        DidOpenTextDocumentParams().apply {
+                            textDocument = TextDocumentItem().apply {
+                                this.uri = uri
+                                text = file.inputStream.readAllBytes().decodeToString()
+                                languageId = file.fileType.name.lowercase()
+                                version = file.modificationStamp.toInt()
+                            }
                         }
-                    }
-                )
+                    )
+                }
             }
         }
     }
@@ -81,14 +97,16 @@ class TextDocumentServiceHandler(
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             val file = FileDocumentManager.getInstance().getFile(document) ?: return@executeIfRunning
             toUriString(file)?.let { uri ->
-                languageServer.textDocumentService.didSave(
-                    DidSaveTextDocumentParams().apply {
-                        textDocument = TextDocumentIdentifier().apply {
-                            this.uri = uri
+                pluginAwareExecuteOnPooledThread {
+                    languageServer.textDocumentService.didSave(
+                        DidSaveTextDocumentParams().apply {
+                            textDocument = TextDocumentIdentifier().apply {
+                                this.uri = uri
+                            }
+                            text = document.text
                         }
-                        text = document.text
-                    }
-                )
+                    )
+                }
             }
         }
     }
@@ -141,4 +159,28 @@ class TextDocumentServiceHandler(
             }
         }
     }
+
+    private fun realTimeEdit(event: DocumentEvent) {
+        AmazonQLspService.executeIfRunning(project) { languageServer ->
+            pluginAwareExecuteOnPooledThread {
+                val vFile = FileDocumentManager.getInstance().getFile(event.document) ?: return@pluginAwareExecuteOnPooledThread
+                toUriString(vFile)?.let { uri ->
+                    languageServer.textDocumentService.didChange(
+                        DidChangeTextDocumentParams().apply {
+                            textDocument = VersionedTextDocumentIdentifier().apply {
+                                this.uri = uri
+                                version = event.document.modificationStamp.toInt()
+                            }
+                            contentChanges = listOf(
+                                TextDocumentContentChangeEvent().apply {
+                                    text = event.document.text
+                                }
+                            )
+                        }
+                    )
+                }
+            }
+        }
+        // Process document changes here
+    }
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandlerTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/textdocument/TextDocumentServiceHandlerTest.kt
@@ -3,29 +3,26 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.textdocument
 
-import com.intellij.openapi.Disposable
-import com.intellij.openapi.application.Application
-import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.components.serviceIfCreated
+import com.intellij.openapi.application.writeAction
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileDocumentManager
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileTypes.FileType
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.newvfs.events.VFileContentChangeEvent
 import com.intellij.openapi.vfs.newvfs.events.VFileEvent
-import com.intellij.util.messages.MessageBus
-import com.intellij.util.messages.MessageBusConnection
+import com.intellij.openapi.vfs.writeText
+import com.intellij.testFramework.DisposableRule
+import com.intellij.testFramework.fixtures.CodeInsightTestFixture
+import com.intellij.testFramework.fixtures.IdeaTestFixtureFactory
+import com.intellij.testFramework.replaceService
 import io.mockk.every
-import io.mockk.just
 import io.mockk.mockk
 import io.mockk.mockkObject
 import io.mockk.mockkStatic
-import io.mockk.runs
 import io.mockk.slot
 import io.mockk.verify
 import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.withContext
 import org.assertj.core.api.Assertions.assertThat
 import org.eclipse.lsp4j.DidChangeTextDocumentParams
 import org.eclipse.lsp4j.DidCloseTextDocumentParams
@@ -34,42 +31,51 @@ import org.eclipse.lsp4j.DidSaveTextDocumentParams
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.eclipse.lsp4j.services.TextDocumentService
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
+import software.aws.toolkits.jetbrains.core.coroutines.EDT
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.FileUriUtil
+import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
+import software.aws.toolkits.jetbrains.utils.satisfiesKt
 import java.net.URI
 import java.nio.file.Path
-import java.util.concurrent.Callable
 import java.util.concurrent.CompletableFuture
+import kotlin.collections.first
 
 class TextDocumentServiceHandlerTest {
-    private lateinit var project: Project
-    private lateinit var mockFileEditorManager: FileEditorManager
     private lateinit var mockLanguageServer: AmazonQLanguageServer
     private lateinit var mockTextDocumentService: TextDocumentService
     private lateinit var sut: TextDocumentServiceHandler
-    private lateinit var mockApplication: Application
+
+    @get:Rule
+    val projectRule = object : CodeInsightTestFixtureRule() {
+        override fun createTestFixture(): CodeInsightTestFixture {
+            val fixtureFactory = IdeaTestFixtureFactory.getFixtureFactory()
+            val fixtureBuilder = fixtureFactory.createLightFixtureBuilder(testDescription, testName)
+            val newFixture = fixtureFactory
+                .createCodeInsightFixture(fixtureBuilder.fixture, fixtureFactory.createTempDirTestFixture())
+            newFixture.setUp()
+            newFixture.testDataPath = testDataPath
+
+            return newFixture
+        }
+    }
+
+    @get:Rule
+    val disposableRule = DisposableRule()
 
     @Before
     fun setup() {
-        project = mockk<Project>()
         mockTextDocumentService = mockk<TextDocumentService>()
         mockLanguageServer = mockk<AmazonQLanguageServer>()
 
-        mockApplication = mockk<Application>()
-        mockkStatic(ApplicationManager::class)
-        every { ApplicationManager.getApplication() } returns mockApplication
-        every { mockApplication.executeOnPooledThread(any<Callable<*>>()) } answers {
-            CompletableFuture.completedFuture(firstArg<Callable<*>>().call())
-        }
-
         // Mock the LSP service
-        val mockLspService = mockk<AmazonQLspService>()
+        val mockLspService = mockk<AmazonQLspService>(relaxed = true)
 
         // Mock the service methods on Project
-        every { project.getService(AmazonQLspService::class.java) } returns mockLspService
-        every { project.serviceIfCreated<AmazonQLspService>() } returns mockLspService
+        projectRule.project.replaceService(AmazonQLspService::class.java, mockLspService, disposableRule.disposable)
 
         // Mock the LSP service's executeSync method as a suspend function
         every {
@@ -86,19 +92,7 @@ class TextDocumentServiceHandlerTest {
         every { mockTextDocumentService.didOpen(any()) } returns Unit
         every { mockTextDocumentService.didClose(any()) } returns Unit
 
-        // Mock message bus
-        val messageBus = mockk<MessageBus>()
-        every { project.messageBus } returns messageBus
-        val mockConnection = mockk<MessageBusConnection>()
-        every { messageBus.connect(any<Disposable>()) } returns mockConnection
-        every { mockConnection.subscribe(any(), any()) } just runs
-
-        // Mock FileEditorManager
-        mockFileEditorManager = mockk<FileEditorManager>()
-        every { mockFileEditorManager.openFiles } returns emptyArray()
-        every { project.getService(FileEditorManager::class.java) } returns mockFileEditorManager
-
-        sut = TextDocumentServiceHandler(project, mockk())
+        sut = TextDocumentServiceHandler(projectRule.project, mockk())
     }
 
     @Test
@@ -136,41 +130,39 @@ class TextDocumentServiceHandlerTest {
 
     @Test
     fun `didOpen runs on service init`() = runTest {
-        val uri = URI.create("file:///test/path/file.txt")
         val content = "test content"
-        val file = createMockVirtualFile(uri, content)
-
-        every { mockFileEditorManager.openFiles } returns arrayOf(file)
+        val file = withContext(EDT) {
+            projectRule.fixture.createFile("name", content).also { projectRule.fixture.openFileInEditor(it) }
+        }
 
-        sut = TextDocumentServiceHandler(project, mockk())
+        sut = TextDocumentServiceHandler(projectRule.project, mockk())
 
-        val paramsSlot = slot<DidOpenTextDocumentParams>()
+        val paramsSlot = mutableListOf<DidOpenTextDocumentParams>()
         verify { mockTextDocumentService.didOpen(capture(paramsSlot)) }
 
-        with(paramsSlot.captured.textDocument) {
-            assertThat(this.uri).isEqualTo(normalizeFileUri(uri.toString()))
-            assertThat(text).isEqualTo(content)
-            assertThat(languageId).isEqualTo("java")
-            assertThat(version).isEqualTo(1)
+        assertThat(paramsSlot.first().textDocument).satisfiesKt {
+            assertThat(it.uri).isEqualTo(file.toNioPath().toUri().toString())
+            assertThat(it.text).isEqualTo(content)
+            assertThat(it.languageId).isEqualTo("plain_text")
         }
     }
 
     @Test
     fun `didOpen runs on fileOpened`() = runTest {
-        val uri = URI.create("file:///test/path/file.txt")
         val content = "test content"
-        val file = createMockVirtualFile(uri, content)
+        val file = withContext(EDT) {
+            projectRule.fixture.createFile("name", content).also { projectRule.fixture.openFileInEditor(it) }
+        }
 
         sut.fileOpened(mockk(), file)
 
-        val paramsSlot = slot<DidOpenTextDocumentParams>()
+        val paramsSlot = mutableListOf<DidOpenTextDocumentParams>()
         verify { mockTextDocumentService.didOpen(capture(paramsSlot)) }
 
-        with(paramsSlot.captured.textDocument) {
-            assertThat(this.uri).isEqualTo(normalizeFileUri(uri.toString()))
-            assertThat(text).isEqualTo(content)
-            assertThat(languageId).isEqualTo("java")
-            assertThat(version).isEqualTo(1)
+        assertThat(paramsSlot.first().textDocument).satisfiesKt {
+            assertThat(it.uri).isEqualTo(file.toNioPath().toUri().toString())
+            assertThat(it.text).isEqualTo(content)
+            assertThat(it.languageId).isEqualTo("plain_text")
         }
     }
 
@@ -189,38 +181,23 @@ class TextDocumentServiceHandlerTest {
 
     @Test
     fun `didChange runs on content change events`() = runTest {
-        val uri = URI.create("file:///test/path/file.txt")
-        val document = mockk<Document> {
-            every { text } returns "changed content"
-            every { modificationStamp } returns 123L
-        }
-
-        val file = createMockVirtualFile(uri)
-
-        val changeEvent = mockk<VFileContentChangeEvent> {
-            every { this@mockk.file } returns file
-        }
-
-        // Mock FileDocumentManager
-        val fileDocumentManager = mockk<FileDocumentManager> {
-            every { getCachedDocument(file) } returns document
-        }
+        val file = withContext(EDT) {
+            projectRule.fixture.createFile("name", "").also {
+                projectRule.fixture.openFileInEditor(it)
 
-        mockkStatic(FileDocumentManager::class) {
-            every { FileDocumentManager.getInstance() } returns fileDocumentManager
-
-            // Call the handler method
-            sut.after(mutableListOf(changeEvent))
+                writeAction {
+                    it.writeText("changed content")
+                }
+            }
         }
 
         // Verify the correct LSP method was called with matching parameters
-        val paramsSlot = slot<DidChangeTextDocumentParams>()
+        val paramsSlot = mutableListOf<DidChangeTextDocumentParams>()
         verify { mockTextDocumentService.didChange(capture(paramsSlot)) }
 
-        with(paramsSlot.captured) {
-            assertThat(textDocument.uri).isEqualTo(normalizeFileUri(uri.toString()))
-            assertThat(textDocument.version).isEqualTo(123)
-            assertThat(contentChanges[0].text).isEqualTo("changed content")
+        assertThat(paramsSlot.first()).satisfiesKt {
+            assertThat(it.textDocument.uri).isEqualTo(file.toNioPath().toUri().toString())
+            assertThat(it.contentChanges[0].text).isEqualTo("changed content")
         }
     }
 
@@ -335,6 +312,11 @@ class TextDocumentServiceHandlerTest {
             return uri
         }
 
+        if (uri.startsWith("file://C:/")) {
+            val path = uri.substringAfter("file://C:/")
+            return "file:///C:/$path"
+        }
+
         val path = uri.substringAfter("file:///")
         return "file:///C:/$path"
     }

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +{
 +  "type" : "bugfix",
 +  "description" : "Fix UI freezes that may occur when interacting with large files in the editor"
 +}