Fix logic for picking-up selected editor when using remote (#2349)

Fixes CODY-3831

## Changes

`FileEditorManager.getInstance(it).selectedTextEditor` does not report
selected editor properly in the remote mode.
Instead we should always use
`FileEditorManager.getInstance(project).selectedTextEditorWithRemotes`
or pick proper editor from the list of all open ones.

## Test plan

1. Start IntelliJ in the dev remote mode (by using `:customRunIde
-PsplitMode=true`)
2. Run any edit command
3. Do any focus change in the editor - there should be no desync issues
reported in the Cody log window

src/main/java/com/sourcegraph/website/FileActionBase.java

@@ -5,13 +5,13 @@
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.editor.*;
 import com.intellij.openapi.fileEditor.FileDocumentManager;
-import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.sourcegraph.common.ErrorNotification;
 import com.sourcegraph.common.ui.DumbAwareEDTAction;
 import com.sourcegraph.find.PreviewContent;
 import com.sourcegraph.find.SourcegraphVirtualFile;
+import com.sourcegraph.utils.CodyEditorUtil;
 import com.sourcegraph.vcs.RepoInfo;
 import com.sourcegraph.vcs.RepoUtil;
 import com.sourcegraph.vcs.VCSType;
@@ -28,7 +28,7 @@ public void actionPerformed(@NotNull AnActionEvent event) {
     if (project == null) {
       return;
     }
-    Editor editor = FileEditorManager.getInstance(project).getSelectedTextEditor();
+    Editor editor = CodyEditorUtil.getFirstSelectedEditor(project);
     if (editor == null) {
       return;
     }

src/main/java/com/sourcegraph/website/SearchActionBase.java

@@ -7,13 +7,13 @@
 import com.intellij.openapi.editor.Editor;
 import com.intellij.openapi.editor.SelectionModel;
 import com.intellij.openapi.fileEditor.FileDocumentManager;
-import com.intellij.openapi.fileEditor.FileEditorManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.sourcegraph.common.BrowserOpener;
 import com.sourcegraph.common.ui.DumbAwareEDTAction;
 import com.sourcegraph.find.SourcegraphVirtualFile;
+import com.sourcegraph.utils.CodyEditorUtil;
 import com.sourcegraph.vcs.RepoInfo;
 import com.sourcegraph.vcs.RepoUtil;
 import com.sourcegraph.vcs.VCSType;
@@ -31,10 +31,13 @@ public void actionPerformedMode(@NotNull AnActionEvent event, @NotNull Scope sco
     if (selectedText == null || selectedText.length() == 0) {
       return;
     }
+    Editor editor = CodyEditorUtil.getFirstSelectedEditor(project);
+    if (editor == null) {
+      return;
+    }
+
     //noinspection ConstantConditions selectedText != null, so the editor can't be null.
-    VirtualFile currentFile =
-        FileDocumentManager.getInstance()
-            .getFile(FileEditorManager.getInstance(project).getSelectedTextEditor().getDocument());
+    VirtualFile currentFile = FileDocumentManager.getInstance().getFile(editor.getDocument());
     assert currentFile != null; // selectedText != null, so this can't be null.
 
     if (currentFile instanceof SourcegraphVirtualFile) {
@@ -103,7 +106,7 @@ private String getSelectedText(@Nullable Project project) {
       return null;
     }
 
-    Editor editor = FileEditorManager.getInstance(project).getSelectedTextEditor();
+    Editor editor = CodyEditorUtil.getFirstSelectedEditor(project);
     if (editor == null) {
       return null;
     }

src/main/kotlin/com/sourcegraph/cody/chat/actions/BaseCommandAction.kt

@@ -5,7 +5,6 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.application.ReadAction
 import com.intellij.openapi.fileEditor.FileDocumentManager
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.util.concurrency.AppExecutorUtil
 import com.sourcegraph.cody.agent.CodyAgentService
@@ -16,6 +15,7 @@ import com.sourcegraph.cody.ignore.ActionInIgnoredFileNotification
 import com.sourcegraph.cody.ignore.IgnoreOracle
 import com.sourcegraph.cody.ignore.IgnorePolicy
 import com.sourcegraph.common.ui.DumbAwareEDTAction
+import com.sourcegraph.utils.CodyEditorUtil
 import java.util.concurrent.Callable
 
 abstract class BaseCommandAction : DumbAwareEDTAction() {
@@ -28,7 +28,7 @@ abstract class BaseCommandAction : DumbAwareEDTAction() {
 
   open fun doAction(project: Project) {
     ApplicationManager.getApplication().assertIsDispatchThread()
-    FileEditorManager.getInstance(project).selectedTextEditor?.let { editor ->
+    CodyEditorUtil.getSelectedEditors(project).firstOrNull()?.let { editor ->
       val file = FileDocumentManager.getInstance().getFile(editor.document)
       val protocolFile =
           file?.let { ProtocolTextDocument.fromVirtualEditorFile(editor, it) } ?: return

src/main/kotlin/com/sourcegraph/cody/edit/lenses/EditCodeVisionProvider.kt

@@ -9,6 +9,7 @@ import com.intellij.codeInsight.codeVision.ui.model.ClickableRichTextCodeVisionE
 import com.intellij.codeInsight.codeVision.ui.model.richText.RichText
 import com.intellij.icons.AllIcons
 import com.intellij.openapi.actionSystem.ActionManager
+import com.intellij.openapi.actionSystem.ActionPlaces
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.actionSystem.PlatformDataKeys
@@ -126,7 +127,7 @@ abstract class EditCodeVisionProvider(private val metadata: EditCodeVisionProvid
           AnActionEvent(
               event,
               dataContext,
-              "",
+              ActionPlaces.EDITOR_INLAY,
               action.templatePresentation.clone(),
               ActionManager.getInstance(),
               0)

src/main/kotlin/com/sourcegraph/cody/edit/lenses/LensesService.kt

@@ -3,14 +3,11 @@ package com.sourcegraph.cody.edit.lenses
 import com.intellij.codeInsight.codeVision.CodeVisionHost
 import com.intellij.codeInsight.codeVision.CodeVisionInitializer
 import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.psi.PsiDocumentManager
 import com.sourcegraph.cody.agent.protocol_generated.ProtocolCodeLens
 import com.sourcegraph.config.ConfigUtil
 import com.sourcegraph.utils.CodyEditorUtil
@@ -66,22 +63,19 @@ class LensesService(val project: Project) {
 
     runInEdt {
       if (project.isDisposed) return@runInEdt
-      val editor = FileEditorManager.getInstance(project).selectedTextEditor
-      CodeVisionInitializer.getInstance(project)
-          .getCodeVisionHost()
-          .invalidateProvider(
-              CodeVisionHost.LensInvalidateSignal(
-                  editor, EditCodeVisionProvider.allEditProviders().map { it.id }))
+      CodyEditorUtil.getSelectedEditors(project).forEach { editor ->
+        CodeVisionInitializer.getInstance(project)
+            .getCodeVisionHost()
+            .invalidateProvider(
+                CodeVisionHost.LensInvalidateSignal(
+                    editor, EditCodeVisionProvider.allEditProviders().map { it.id }))
+      }
     }
     listeners.forEach { it.onLensesUpdate(vf, codeLens) }
   }
 
   fun getLenses(editor: Editor): List<ProtocolCodeLens> {
-    val document = editor.document
-    val file =
-        runReadAction { PsiDocumentManager.getInstance(project).getPsiFile(document) }
-            ?: return emptyList()
-    val vf = file.viewProvider.virtualFile
+    val vf = editor.virtualFile
 
     synchronized(this) {
       return lensGroups[vf] ?: emptyList()

src/main/kotlin/com/sourcegraph/cody/editor/CommandListener.kt

@@ -2,18 +2,17 @@ package com.sourcegraph.cody.editor
 
 import com.intellij.openapi.command.CommandEvent
 import com.intellij.openapi.command.CommandListener
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.autocomplete.CodyAutocompleteManager
+import com.sourcegraph.utils.CodyEditorUtil
 import com.sourcegraph.utils.CodyEditorUtil.VIM_EXIT_INSERT_MODE_ACTION
 
 class CodyCommandListener(val project: Project) : CommandListener {
   override fun commandFinished(event: CommandEvent) {
     if (event.commandName.isNullOrBlank() ||
         event.commandName.equals(VIM_EXIT_INSERT_MODE_ACTION)) {
-      val fileEditorManager = FileEditorManager.getInstance(this.project)
-      fileEditorManager.selectedTextEditor?.let {
-        CodyAutocompleteManager.instance.disposeInlays(it)
+      CodyEditorUtil.getSelectedEditors(project).forEach { editor ->
+        CodyAutocompleteManager.instance.disposeInlays(editor)
       }
     }
   }

src/main/kotlin/com/sourcegraph/cody/ignore/IgnoreOracle.kt

@@ -7,14 +7,14 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.fileEditor.FileDocumentManager
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.intellij.util.containers.SLRUMap
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.protocol.IgnoreTestParams
 import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
 import com.sourcegraph.cody.statusbar.CodyStatusService
+import com.sourcegraph.utils.CodyEditorUtil
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.TimeoutException
@@ -44,11 +44,12 @@ class IgnoreOracle(private val project: Project) {
     // Synthesize a focus event for the current editor, if any,
     // to fetch and cache ignore state for it.
     runInEdt {
-      val editor = FileEditorManager.getInstance(project).selectedTextEditor
-      if (willFocusUri == null && editor != null) {
-        val uri = ProtocolTextDocument.fromEditor(editor)?.uri
-        if (uri != null) {
-          focusedFileDidChange(uri)
+      CodyEditorUtil.getSelectedEditors(project).forEach { editor ->
+        if (willFocusUri == null) {
+          val uri = ProtocolTextDocument.fromEditor(editor)?.uri
+          if (uri != null) {
+            focusedFileDidChange(uri)
+          }
         }
       }
     }

src/main/kotlin/com/sourcegraph/cody/listeners/CodyDocumentListener.kt

@@ -3,7 +3,6 @@ package com.sourcegraph.cody.listeners
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.editor.event.BulkAwareDocumentListener
 import com.intellij.openapi.editor.event.DocumentEvent
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.protocol.CompletionItemParams
@@ -13,6 +12,7 @@ import com.sourcegraph.cody.autocomplete.action.AcceptCodyAutocompleteAction
 import com.sourcegraph.cody.chat.CodeEditorFactory
 import com.sourcegraph.cody.telemetry.TelemetryV2
 import com.sourcegraph.cody.vscode.InlineCompletionTriggerKind
+import com.sourcegraph.utils.CodyEditorUtil
 
 class CodyDocumentListener(val project: Project) : BulkAwareDocumentListener {
 
@@ -37,10 +37,7 @@ class CodyDocumentListener(val project: Project) : BulkAwareDocumentListener {
   }
 
   private fun handleDocumentEvent(event: DocumentEvent) {
-    val editor = FileEditorManager.getInstance(project).selectedTextEditor
-    if (editor?.document != event.document) {
-      return
-    }
+    val editor = CodyEditorUtil.getEditorForDocument(event.document) ?: return
 
     logCodeCopyPastedFromChat(event)
     CodyAutocompleteManager.instance.clearAutocompleteSuggestions(editor)

src/main/kotlin/com/sourcegraph/cody/listeners/CodyFileEditorListener.kt

@@ -2,7 +2,6 @@ package com.sourcegraph.cody.listeners
 
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.diagnostic.Logger
-import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
@@ -13,6 +12,7 @@ import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.agent.CodyAgentService.Companion.withAgent
 import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
 import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument.Companion.fromVirtualEditorFile
+import com.sourcegraph.utils.CodyEditorUtil
 
 class CodyFileEditorListener : FileEditorManagerListener {
   private val logger = Logger.getInstance(CodyFileEditorListener::class.java)
@@ -54,7 +54,7 @@ class CodyFileEditorListener : FileEditorManagerListener {
       ApplicationManager.getApplication().invokeLater {
         val fileDocumentManager = FileDocumentManager.getInstance()
 
-        EditorFactory.getInstance().allEditors.forEach { editor ->
+        CodyEditorUtil.getAllOpenEditors().forEach { editor ->
           fileDocumentManager.getFile(editor.document)?.let { file ->
             try {
               val textDocument = fromVirtualEditorFile(editor, file)
@@ -66,7 +66,7 @@ class CodyFileEditorListener : FileEditorManagerListener {
         }
 
         if (project.isDisposed) return@invokeLater
-        FileEditorManager.getInstance(project).selectedTextEditor?.let { editor ->
+        CodyEditorUtil.getSelectedEditors(project).forEach { editor ->
           val file = fileDocumentManager.getFile(editor.document)
           try {
             val textDocument = fromVirtualEditorFile(editor, file!!)

src/main/kotlin/com/sourcegraph/cody/ui/web/WebUIHost.kt

@@ -8,14 +8,14 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.actionSystem.impl.SimpleDataContext
 import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.fileEditor.FileEditorManager
 import com.intellij.openapi.options.ShowSettingsUtil
 import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.agent.*
 import com.sourcegraph.cody.agent.protocol.WebviewOptions
 import com.sourcegraph.cody.config.CodyAuthenticationManager
 import com.sourcegraph.cody.config.ui.AccountConfigurable
 import com.sourcegraph.cody.config.ui.CodyConfigurable
+import com.sourcegraph.utils.CodyEditorUtil
 import java.net.URLDecoder
 
 // TODO:
@@ -87,7 +87,7 @@ internal class WebUIHostImpl(
         val actionManager = ActionManager.getInstance()
         val action = actionManager.getAction("cody.editCodeAction")
         val dataContext =
-            FileEditorManager.getInstance(project).selectedTextEditor?.let { editor ->
+            CodyEditorUtil.getSelectedEditors(project).firstOrNull()?.let { editor ->
               SimpleDataContext.getSimpleContext(CommonDataKeys.EDITOR, editor)
             } ?: SimpleDataContext.EMPTY_CONTEXT
 

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

@@ -75,13 +75,24 @@ object CodyEditorUtil {
   }
 
   @JvmStatic
-  fun getFocusedEditorForAnActionEvent(e: AnActionEvent): Editor? {
-    return e.project?.let { FileEditorManager.getInstance(it).selectedTextEditor }
+  fun getSelectedEditors(project: Project): Array<out Editor> {
+    return FileEditorManager.getInstance(project).selectedTextEditorWithRemotes
+  }
+
+  @JvmStatic
+  fun getFirstSelectedEditor(project: Project): Editor? {
+    return getSelectedEditors(project).firstOrNull()
+  }
+
+  @JvmStatic
+  fun getEditorForDocument(document: Document): Editor? {
+    return getAllOpenEditors().find { it.document == document }
   }
 
   @JvmStatic
   fun getLanguageForFocusedEditor(e: AnActionEvent): Language? {
-    return getFocusedEditorForAnActionEvent(e)?.let { getLanguage(it) }
+    val project = e.project ?: return null
+    return getSelectedEditors(project).firstOrNull()?.let { getLanguage(it) }
   }
 
   @JvmStatic