Fix IndexOutOfBounds from the highlight pass (#2252)

## Test plan
Try these examples:
- #2110
- #1988

src/main/kotlin/com/sourcegraph/cody/agent/intellij_extensions/Document.kt

@@ -13,6 +13,14 @@ fun Document.codyPosition(offset: Int): Position {
 }
 
 fun Document.codyRange(startOffset: Int, endOffset: Int): Range {
+  if (startOffset < 0 ||
+      startOffset > this.textLength ||
+      endOffset > this.textLength ||
+      startOffset > endOffset) {
+    throw IllegalArgumentException(
+        "codyRange error - startOffset: $startOffset, endOffset: $endOffset, textLength: ${this.textLength}")
+  }
+
   val startLine = this.getLineNumber(startOffset)
   val lineStartOffset1 = this.getLineStartOffset(startLine)
   val startCharacter = startOffset - lineStartOffset1

src/main/kotlin/com/sourcegraph/cody/inspections/CodyFixHighlightPass.kt

@@ -9,26 +9,23 @@ import com.intellij.codeInsight.daemon.DaemonCodeAnalyzer
 import com.intellij.codeInsight.daemon.impl.DaemonCodeAnalyzerImpl
 import com.intellij.codeInsight.daemon.impl.HighlightInfo
 import com.intellij.lang.annotation.HighlightSeverity
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.util.ProgressIndicatorUtils
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
-import com.intellij.util.concurrency.AppExecutorUtil
+import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.intellij_extensions.codyRange
-import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
+import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument.Companion.uriFor
 import com.sourcegraph.cody.agent.protocol_generated.CodeActions_ProvideParams
 import com.sourcegraph.cody.agent.protocol_generated.Diagnostics_PublishParams
 import com.sourcegraph.cody.agent.protocol_generated.ProtocolDiagnostic
 import com.sourcegraph.cody.agent.protocol_generated.ProtocolLocation
 import com.sourcegraph.cody.agent.protocol_generated.Range
+import com.sourcegraph.utils.ThreadingUtil
 import java.util.concurrent.CompletableFuture
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.TimeoutException
-import java.util.concurrent.atomic.AtomicReference
 
 class CodyFixHighlightPass(val file: PsiFile, val editor: Editor) :
     TextEditorHighlightingPass(file.project, editor.document, false) {
@@ -37,84 +34,82 @@ class CodyFixHighlightPass(val file: PsiFile, val editor: Editor) :
   private var myHighlights = emptyList<HighlightInfo>()
   private val myRangeActions = mutableMapOf<Range, List<CodeActionQuickFixParams>>()
 
-  // We are invoked by the superclass on a daemon/worker thread, but we need the EDT
-  // for this and perhaps other things (e.g. Document.codyRange). So we set things up
-  // to block the caller and fetch what we need on the EDT.
-  private val protocolTextDocumentFuture: CompletableFuture<ProtocolTextDocument> =
-      CompletableFuture.supplyAsync(
-          {
-            val result = AtomicReference<ProtocolTextDocument>()
-            ApplicationManager.getApplication().invokeAndWait {
-              result.set(ProtocolTextDocument.fromVirtualFile(file.virtualFile))
-            }
-            result.get()
-          },
-          AppExecutorUtil.getAppExecutorService())
-
   override fun doCollectInformation(progress: ProgressIndicator) {
-    // Fetch the protocol text document on the EDT.
-    val protocolTextDocument =
-        try {
-          protocolTextDocumentFuture.get(5, TimeUnit.SECONDS)
-        } catch (e: TimeoutException) {
-          logger.warn("Timeout fetching protocol text document")
-          return
-        }
-
     if (!DaemonCodeAnalyzer.getInstance(file.project).isHighlightingAvailable(file) ||
         progress.isCanceled) {
       // wait until after code-analysis is completed
       return
     }
+    val uri = uriFor(file.virtualFile)
 
     myRangeActions.clear()
 
     myHighlights =
         DaemonCodeAnalyzerImpl.getHighlights(editor.document, HighlightSeverity.ERROR, file.project)
 
     val protocolDiagnostics =
-        myHighlights
-            // TODO: We need to check how Enum comparison works to check if we can do things like >=
-            // HighlightSeverity.INFO
-            .filter { it.severity == HighlightSeverity.ERROR }
-            .mapNotNull {
-              try {
-                ProtocolDiagnostic(
-                    message = it.description,
-                    severity =
-                        "error", // TODO: Wait for CODY-2882. This isn't currently used by the
-                    // agent,
-                    // so we just keep our lives simple.
-                    location =
-                        // TODO: Rik Nauta -- Got incorrect range; see QA report Aug 6 2024.
-                        ProtocolLocation(
-                            uri = protocolTextDocument.uri,
-                            range = document.codyRange(it.startOffset, it.endOffset)),
-                    code = it.problemGroup?.problemName)
-              } catch (x: Exception) {
-                // Don't allow range errors to throw user-visible exceptions (QA found this).
-                logger.warn("Failed to convert highlight to protocol diagnostic", x)
-                null
+        ThreadingUtil.runInEdtAndGet {
+          myHighlights
+              // TODO: We need to check how Enum comparison works to check if we can do things like
+              // >=
+              // HighlightSeverity.INFO
+              .asSequence()
+              .filter { it.severity == HighlightSeverity.ERROR }
+              .filter { it.startOffset <= document.textLength }
+              .filter { it.endOffset <= document.textLength }
+              .filter { it.startOffset <= it.endOffset }
+              .mapNotNull {
+                try {
+                  ProtocolDiagnostic(
+                      message = it.description,
+                      severity =
+                          "error", // TODO: Wait for CODY-2882. This isn't currently used by the
+                      // agent,
+                      // so we just keep our lives simple.
+                      location =
+                          // TODO: Rik Nauta -- Got incorrect range; see QA report Aug 6 2024.
+                          ProtocolLocation(
+                              uri = uri, range = document.codyRange(it.startOffset, it.endOffset)),
+                      code = it.problemGroup?.problemName)
+                } catch (x: Exception) {
+                  // Don't allow range errors to throw user-visible exceptions (QA found this).
+                  logger.warn("Failed to convert highlight to protocol diagnostic", x)
+                  null
+                }
               }
-            }
+              .toList()
+        }
+
+    if (protocolDiagnostics.isEmpty()) {
+      return
+    }
+
     val done = CompletableFuture<Unit>()
     CodyAgentService.withAgentRestartIfNeeded(file.project) { agent ->
       try {
         agent.server.diagnostics_publish(
             Diagnostics_PublishParams(diagnostics = protocolDiagnostics))
+
         for (highlight in myHighlights) {
           if (progress.isCanceled) {
             break
           }
-          val range = document.codyRange(highlight.startOffset, highlight.endOffset)
+
+          val range =
+              ThreadingUtil.runInEdtAndGet {
+                if (highlight.startOffset > document.textLength ||
+                    highlight.endOffset > document.textLength ||
+                    highlight.startOffset > highlight.endOffset) {
+                  return@runInEdtAndGet null
+                }
+
+                return@runInEdtAndGet document.codyRange(highlight.startOffset, highlight.endOffset)
+              } ?: break
 
           if (myRangeActions.containsKey(range)) {
             continue
           }
-          val location =
-              ProtocolLocation(
-                  uri = protocolTextDocument.uri,
-                  range = document.codyRange(highlight.startOffset, highlight.endOffset))
+          val location = ProtocolLocation(uri = uri, range = range)
           val provideResponse =
               agent.server
                   .codeActions_provide(
@@ -133,10 +128,17 @@ class CodyFixHighlightPass(val file: PsiFile, val editor: Editor) :
     ProgressIndicatorUtils.awaitWithCheckCanceled(done, progress)
   }
 
+  @RequiresEdt
   override fun doApplyInformationToEditor() {
     for (highlight in myHighlights) {
       highlight.unregisterQuickFix { it.familyName == CodeActionQuickFix.FAMILY_NAME }
 
+      if (highlight.startOffset > document.textLength ||
+          highlight.endOffset > document.textLength ||
+          highlight.startOffset > highlight.endOffset) {
+        break
+      }
+
       val range = document.codyRange(highlight.startOffset, highlight.endOffset)
       for (action in myRangeActions[range].orEmpty()) {
         highlight.registerFix(

src/main/kotlin/com/sourcegraph/utils/CodyFormatter.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.codeStyle.CodeStyleManager
+import kotlin.math.max
 import kotlin.math.min
 
 class CodyFormatter {
@@ -35,9 +36,8 @@ class CodyFormatter {
                 .createFileFromText("TEMP", file.fileType, appendedString)
 
         val codeStyleManager = CodeStyleManager.getInstance(project)
-        // This will fail if cursor < range.startOffset + completionText.length
-        // TODO change the signature of this method or at least validate the arguments
-        codeStyleManager.reformatText(psiFile, cursor, range.startOffset + completionText.length)
+        val endOffset = max(cursor, range.startOffset + completionText.length)
+        codeStyleManager.reformatText(psiFile, cursor, endOffset)
 
         // Fix for the IJ formatting bug which removes spaces even before the given formatting
         // range.

src/main/kotlin/com/sourcegraph/utils/ThreadingUtil.kt

@@ -0,0 +1,19 @@
+package com.sourcegraph.utils
+
+import java.util.concurrent.CompletableFuture
+import javax.swing.SwingUtilities.invokeAndWait
+
+object ThreadingUtil {
+
+  fun <T> runInEdtAndGet(task: () -> T): T {
+    val future = CompletableFuture<T>()
+    invokeAndWait {
+      try {
+        future.complete(task())
+      } catch (exception: Exception) {
+        future.completeExceptionally(exception)
+      }
+    }
+    return future.get()
+  }
+}