Initial implementation for Inline Edit and Document Code (#440)

Committing before it is ready because Olaf would like to play with it.
Current state: Document Code makes a successful roundtrip to the agent.
And we have Edit Instructions dialog. The rest is not wired up yet.

## Test plan

Locally tested, but the code is by no means ready to use yet.

---------

Signed-off-by: Steve Yegge <[email protected]>
Co-authored-by: Dominic Cooney <[email protected]>
Co-authored-by: Beatrix <[email protected]>

.gitignore

@@ -11,6 +11,9 @@
 # IntelliJ project
 *.iml
 
+# User local IDEA run configurations
+.run/
+
 # Build output & caches for IntelliJ plugin development
 build/
 idea-sandbox/

CONTRIBUTING.md

@@ -197,12 +197,15 @@ There are two supported configurations for debugging this way:
     - The Cody extension connects via socket to the "remote" agent
 
 Option 1 is the simplest, and probably makes the most sense for you
-to use if you are uncertain which method to use for debugging.
+to use if you are uncertain which method to use for debugging. Option 2
+is especially useful when you need to set a breakpoint very early in
+the Agent startup.
 
 ## How to set up Run Configurations
 
-Run configurations are basically IDEA's launcher scripts. You will need to create one
-run configuration in each project window, using Run → Edit Configurations.
+Run configurations are basically IDEA's launcher scripts. You will need
+to create one run configuration in each project window, using Run → Edit
+Configurations.
 
 For both debugging setups (Cody-spawns and JB-spawned), you will need:
 
@@ -305,7 +308,7 @@ Create this configuration in the `sourcegraph/cody` project window:
 - Working directory: `~/src/sg/cody` (or wherever you cloned sourcegraph/cody)
 - JavaScript file: `agent/dist/index.js`
     - note: this is a relative path, unlike the run configs for Option 2
-- Application Parameters: <leave empty>
+- Application Parameters: leave empty
 - Environment variables:
     - `CODY_AGENT_DEBUG_REMOTE=true`
     - `CODY_AGENT_DEBUG_PORT=3113`
@@ -352,3 +355,29 @@ To set this up, in the Before launch section of any Run/Debug configuration, do
 You only need to do most of these steps once, when you create the
 External Tool. Then you can use it in any run config that spawns the
 agent, to rebuild it first.
+
+## Debugging VS Code
+
+Sometimes, the TypeScript backend behaves differently when called from
+IntelliJ via Agent than when the same code is called from the VS Code
+extension, and you may need to debug the same code paths through the
+Agent twice. First, when called from the JetBrains extension side, and
+again when called from the VS Code extension side.
+
+To accomplish the latter, you can use VS Code to debug itself. It works
+very similarly to how it works in JetBrains. There are predefined run
+configurations for debugging VS Code Cody in the file
+`.vscode/launch.json` in `sourcegraph/cody`, such as `Launch VS Code
+Extension (Desktop)`.
+
+You do not launch VS Code run configurations from the command palette.
+Instead, use ctrl-shift-D to open the Run and Debug view, and you can
+see the configuration dropdown at the top.
+
+## Known Issues
+
+- Force-stopping the target often corrupts IntelliJ's project indexes,
+  requiring an extra restart of the debugged target to fix them.
+  - Workaround is to exit the target gracefully by quitting each time,
+    using the menus or hotkeys, rather than force-stopping it.
+

src/main/java/com/sourcegraph/cody/CodyFileEditorListener.java

@@ -11,7 +11,7 @@
 import com.sourcegraph.cody.agent.CodyAgent;
 import com.sourcegraph.cody.agent.CodyAgentCodebase;
 import com.sourcegraph.cody.agent.CodyAgentService;
-import com.sourcegraph.cody.agent.protocol.TextDocument;
+import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument;
 import com.sourcegraph.config.ConfigUtil;
 import org.jetbrains.annotations.NotNull;
 
@@ -34,7 +34,8 @@ public void fileClosed(@NotNull FileEditorManager source, @NotNull VirtualFile f
 
     CodyAgentService.withAgent(
         source.getProject(),
-        agent -> agent.getServer().textDocumentDidClose(TextDocument.fromVirtualFile(file)));
+        agent ->
+            agent.getServer().textDocumentDidClose(ProtocolTextDocument.fromVirtualFile(file)));
   }
 
   public static void fileOpened(Project project, CodyAgent codyAgent, @NotNull VirtualFile file) {
@@ -43,7 +44,8 @@ public static void fileOpened(Project project, CodyAgent codyAgent, @NotNull Vir
             .runReadAction(
                 (Computable<Document>) () -> FileDocumentManager.getInstance().getDocument(file));
     if (document != null) {
-      TextDocument textDocument = TextDocument.fromVirtualFile(file, document.getText());
+      ProtocolTextDocument textDocument =
+          ProtocolTextDocument.fromVirtualFile(file, document.getText());
       codyAgent.getServer().textDocumentDidOpen(textDocument);
     }
 

src/main/java/com/sourcegraph/cody/CodyFocusChangeListener.java

@@ -11,7 +11,7 @@
 import com.intellij.openapi.vfs.VirtualFile;
 import com.sourcegraph.cody.agent.CodyAgentCodebase;
 import com.sourcegraph.cody.agent.CodyAgentService;
-import com.sourcegraph.cody.agent.protocol.TextDocument;
+import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument;
 import com.sourcegraph.config.ConfigUtil;
 import org.jetbrains.annotations.NotNull;
 
@@ -54,7 +54,7 @@ public void focusGained(@NotNull Editor editor) {
             Thread.sleep(100);
           } catch (InterruptedException ignored) {
           }
-          agent.getServer().textDocumentDidFocus(TextDocument.fromVirtualFile(file));
+          agent.getServer().textDocumentDidFocus(ProtocolTextDocument.fromVirtualFile(file));
         });
 
     CodyAgentCodebase.getInstance(project).onFileOpened(project, file);

src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java

@@ -3,19 +3,26 @@
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Editor;
-import com.sourcegraph.cody.agent.protocol.DebugMessage;
+import com.sourcegraph.cody.agent.protocol.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification;
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
-/** Implementation of the client part of the Cody agent protocol. */
+/**
+ * Implementation of the client part of the Cody agent protocol. This class dispatches the requests
+ * and notifications sent by the agent.
+ */
 @SuppressWarnings("unused")
 public class CodyAgentClient {
 
   private static final Logger logger = Logger.getInstance(CodyAgentClient.class);
+
+  @Nullable public Editor editor;
+
   // Callback that is invoked when the agent sends a "chat/updateMessageInProgress" notification.
   @Nullable public Consumer<WebviewPostMessageParams> onNewMessage;
 
@@ -26,14 +33,92 @@ public class CodyAgentClient {
   // onSetConfigFeatures
   @Nullable public Consumer<WebviewPostMessageParams> onReceivedWebviewMessage;
 
-  @Nullable public Editor editor;
+  // Callback for the "editTask/didUpdate" notification from the agent.
+  @Nullable private Consumer<EditTask> onEditTaskDidUpdate;
+
+  // Callback for the "editTask/didDelete" notification from the agent.
+  @Nullable private Consumer<EditTask> onEditTaskDidDelete;
+
+  // Callback for the "textDocument/edit" request from the agent.
+  @Nullable private Consumer<TextDocumentEditParams> onTextDocumentEdit;
+
+  // Callback for the "workspace/edit" request from the agent.
+  @Nullable private Consumer<WorkspaceEditParams> onWorkspaceEdit;
+
+  public void setOnEditTaskDidUpdate(@Nullable Consumer<EditTask> callback) {
+    onEditTaskDidUpdate = callback;
+  }
+
+  public void setOnEditTaskDidDelete(@Nullable Consumer<EditTask> callback) {
+    onEditTaskDidDelete = callback;
+  }
+
+  @JsonNotification("editTask/didUpdate")
+  public void editTaskDidUpdate(EditTask params) {
+    onEventThread(
+        () -> {
+          if (onEditTaskDidUpdate != null) {
+            onEditTaskDidUpdate.accept(params);
+          } else {
+            logger.warn("No callback registered for editTask/didUpdate");
+          }
+          return null;
+        });
+  }
+
+  @JsonNotification("editTask/didDelete")
+  public void editTaskDidDelete(EditTask params) {
+    onEventThread(
+        () -> {
+          if (onEditTaskDidDelete != null) {
+            onEditTaskDidDelete.accept(params);
+          } else {
+            logger.warn("No callback registered for editTask/didDelete");
+          }
+          return null;
+        });
+  }
+
+  public void setOnTextDocumentEdit(@Nullable Consumer<TextDocumentEditParams> callback) {
+    onTextDocumentEdit = callback;
+  }
+
+  @JsonRequest("textDocument/edit")
+  public CompletableFuture<Boolean> textDocumentEdit(TextDocumentEditParams params) {
+    return onEventThread(
+        () -> {
+          if (onTextDocumentEdit != null) {
+            onTextDocumentEdit.accept(params);
+          } else {
+            logger.warn("No callback registered for textDocument/edit");
+          }
+          return true;
+        });
+  }
+
+  public void setOnWorkspaceEdit(@Nullable Consumer<WorkspaceEditParams> callback) {
+    onWorkspaceEdit = callback;
+  }
+
+  @JsonRequest("workspace/edit")
+  public CompletableFuture<Boolean> workspaceEdit(WorkspaceEditParams params) {
+    return onEventThread(
+        () -> {
+          if (onWorkspaceEdit != null) {
+            onWorkspaceEdit.accept(params);
+          } else {
+            logger.warn("No callback registered for workspace/edit");
+          }
+          return true;
+        });
+  }
 
   /**
    * Helper to run client request/notification handlers on the IntelliJ event thread. Use this
    * helper for handlers that require access to the IntelliJ editor, for example to read the text
    * contents of the open editor.
    */
-  private <T> CompletableFuture<T> onEventThread(Supplier<T> handler) {
+  private <T> @NotNull CompletableFuture<T> onEventThread(Supplier<T> handler) {
     CompletableFuture<T> result = new CompletableFuture<>();
     ApplicationManager.getApplication()
         .invokeLater(
@@ -47,24 +132,24 @@ private <T> CompletableFuture<T> onEventThread(Supplier<T> handler) {
     return result;
   }
 
+  // Webviews
+  @JsonRequest("webview/create")
+  public CompletableFuture<Void> webviewCreate(WebviewCreateParams params) {
+    logger.error("webview/create This request should not happen if you are using chat/new.");
+    return CompletableFuture.completedFuture(null);
+  }
+
   // =============
   // Notifications
   // =============
 
   @JsonNotification("debug/message")
-  public void debugMessage(DebugMessage msg) {
+  public void debugMessage(@NotNull DebugMessage msg) {
     logger.warn(String.format("%s: %s", msg.getChannel(), msg.getMessage()));
   }
 
-  // Webviews
-  @JsonRequest("webview/create")
-  public CompletableFuture<Void> webviewCreate(WebviewCreateParams params) {
-    logger.error("webview/create This request should not happen if you are using chat/new.");
-    return CompletableFuture.completedFuture(null);
-  }
-
   @JsonNotification("webview/postMessage")
-  public void webviewPostMessage(WebviewPostMessageParams params) {
+  public void webviewPostMessage(@NotNull WebviewPostMessageParams params) {
     ExtensionMessage extensionMessage = params.getMessage();
 
     if (onNewMessage != null

src/main/java/com/sourcegraph/cody/agent/CommandExecuteParams.kt

@@ -1,3 +1,3 @@
 package com.sourcegraph.cody.agent
 
-data class CommandExecuteParams(val command: String, val args: List<String>)
+data class CommandExecuteParams(val command: String, val arguments: List<String>)

src/main/java/com/sourcegraph/cody/vscode/CancellationToken.java

@@ -24,7 +24,7 @@ public void onFinished(Consumer<Boolean> callback) {
           try {
             callback.accept(isCancelled);
           } catch (Exception ignored) {
-            // Do nothing about exceptions in cancelation callbacks
+            // Do nothing about exceptions in cancellation callbacks
           }
         });
   }

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt

@@ -90,8 +90,6 @@ private constructor(
     private const val DEFAULT_AGENT_DEBUG_PORT = 3113 // Also defined in agent/src/cli/jsonrpc.ts
     @JvmField val executorService: ExecutorService = Executors.newCachedThreadPool()
 
-    private fun shouldConnectToDebugAgent() = System.getenv("CODY_AGENT_DEBUG_REMOTE") == "true"
-
     private fun shouldSpawnDebuggableAgent() = System.getenv("CODY_AGENT_DEBUG_INSPECT") == "true"
 
     fun create(project: Project): CompletableFuture<CodyAgent> {
@@ -109,7 +107,10 @@ private constructor(
                       version = ConfigUtil.getPluginVersion(),
                       workspaceRootUri =
                           ConfigUtil.getWorkspaceRootPath(project).toUri().toString(),
-                      extensionConfiguration = ConfigUtil.getAgentConfiguration(project)))
+                      extensionConfiguration = ConfigUtil.getAgentConfiguration(project),
+                      capabilities =
+                          ClientCapabilities(
+                              edit = "enabled", editWorkspace = "enabled", codeLenses = "enabled")))
               .thenApply { info ->
                 logger.info("Connected to Cody agent " + info.name)
                 server.initialized()
@@ -126,7 +127,7 @@ private constructor(
     }
 
     private fun startAgentProcess(): AgentConnection {
-      if (shouldConnectToDebugAgent()) {
+      if (ConfigUtil.shouldConnectToDebugAgent()) {
         return connectToDebugAgent()
       }
       val token = CancellationToken()

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentCodebase.kt

@@ -40,7 +40,7 @@ class CodyAgentCodebase(val project: Project) {
   companion object {
     @JvmStatic
     fun getInstance(project: Project): CodyAgentCodebase {
-      return project.service<CodyAgentCodebase>()
+      return project.service()
     }
   }
 }