Code Actions Support (Ask Cody to Fix) (#1921)

Add the ability for the Agent to provide CodeActions. In practice this means "Ask Cody to Fix" is now supported as a QuickFix on errors shown in JetBrains. https://github.com/user-attachments/assets/f3bad3fd-cd67-427c-baa6-449a4c0dfd48 ## FAQ **Q: What errors are supported?** A: Any errors shown in the JetBrains UI. Cody doesn't do any code analysis, we simply pass the issues that JetBrains is showing along. These Issues (or Inspections as they're known) can come both from JetBrain's language support, external Plugins as well as custom Inspections. **Q: Can Cody explain the error like in VSCode?** A: Not yet, but the foundational support is there now so it will be easy enough to add soon. Fixes CODY-2788, CODY-2830 ## Test plan - Verified that existing unit & integration tests work with Protocol migrated protocol classes - Manually verified that "Ask Cody to Fix" command is correctly shown and invoked. Most of the logic remains on the Agent. - Isolated changes on FixupSession to FixupSessionV2 to not introduce bugs in FixupSessionV1 dependent codepaths which currently have low test coverage.
sourcegraph · Jul 31, 2024 · 1cdb753 · 1cdb753
1 parent 039cebf
commit 1cdb753
Show file tree

Hide file tree

Showing 45 changed files with 486 additions and 305 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -12,6 +12,7 @@ import java.util.EnumSet
 import java.util.jar.JarFile
 import java.util.zip.ZipFile
 import org.jetbrains.changelog.markdownToHTML
+import org.jetbrains.intellij.or
 import org.jetbrains.intellij.tasks.RunPluginVerifierTask.FailureLevel
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import org.jetbrains.kotlin.incremental.deleteDirectoryContents
@@ -212,7 +213,8 @@ fun Test.sharedIntegrationTestConfig(buildCodyDir: File, mode: String) {
 
   val resourcesDir = project.file("src/integrationTest/resources")
   systemProperties(
-      "cody-agent.trace-path" to "$buildDir/sourcegraph/cody-agent-trace.json",
+      "cody-agent.trace-path" to
+          "${layout.buildDirectory.asFile.get()}/sourcegraph/cody-agent-trace.json",
       "cody-agent.directory" to buildCodyDir.parent,
       "sourcegraph.verbose-logging" to "true",
       "cody.autocomplete.enableFormatting" to
@@ -313,7 +315,7 @@ tasks {
     return destination
   }
 
-  val buildCodyDir = buildDir.resolve("sourcegraph").resolve("agent")
+  val buildCodyDir = layout.buildDirectory.asFile.get().resolve("sourcegraph").resolve("agent")
   fun buildCody(): File {
     if (!isForceAgentBuild && (buildCodyDir.listFiles()?.size ?: 0) > 0) {
       println("Cached $buildCodyDir")
@@ -400,7 +402,8 @@ tasks {
   // `./gradlew test` or when testing inside IntelliJ
   val agentProperties =
       mapOf<String, Any>(
-          "cody-agent.trace-path" to "$buildDir/sourcegraph/cody-agent-trace.json",
+          "cody-agent.trace-path" to
+              "${layout.buildDirectory.asFile.get()}/sourcegraph/cody-agent-trace.json",
           "cody-agent.directory" to buildCodyDir.parent,
           "sourcegraph.verbose-logging" to "true",
           "cody-agent.panic-when-out-of-sync" to
@@ -410,11 +413,12 @@ tasks {
           "cody.autocomplete.enableFormatting" to
               (project.property("cody.autocomplete.enableFormatting") ?: "true"))
 
-  fun getIdeaInstallDir(ideaVersion: String): File? {
+  fun getIdeaInstallDir(ideaVersion: String, ideaType: String): File? {
     val gradleHome = project.gradle.gradleUserHomeDir
-    val cacheDir = File(gradleHome, "caches/modules-2/files-2.1/com.jetbrains.intellij.idea/ideaIC")
+    val cacheDir =
+        File(gradleHome, "caches/modules-2/files-2.1/com.jetbrains.intellij.idea/idea$ideaType")
     val ideaDir = File(cacheDir, ideaVersion)
-    return ideaDir.walk().find { it.name == "ideaIC-$ideaVersion" }
+    return ideaDir.walk().find { it.name == "idea$ideaType-$ideaVersion" }
   }
 
   register("copyProtocol") { copyProtocol() }
@@ -496,13 +500,20 @@ tasks {
     agentProperties.forEach { (key, value) -> systemProperty(key, value) }
 
     val platformRuntimeVersion = project.findProperty("platformRuntimeVersion")
-    if (platformRuntimeVersion != null) {
+    val platformRuntimeType = project.findProperty("platformRuntimeType")
+    if (platformRuntimeVersion != null || platformRuntimeType != null) {
       val ideaInstallDir =
-          getIdeaInstallDir(platformRuntimeVersion.toString())
+          getIdeaInstallDir(
+              platformRuntimeVersion.or(project.property("platformVersion")).toString(),
+              platformRuntimeType.or(project.property("platformType")).toString())
               ?: throw GradleException(
-                  "Could not find IntelliJ install for $platformRuntimeVersion")
+                  "Could not find IntelliJ install for version: $platformRuntimeVersion")
       ideDir.set(ideaInstallDir)
     }
+    // TODO: we need to wait to switch to Platform Gradle Plugin 2.0.0 to be able to have separate
+    // runtime plugins
+    // https://github.com/JetBrains/intellij-platform-gradle-plugin/issues/1489
+    // val platformRuntimePlugins = project.findProperty("platformRuntimePlugins")
   }
 
   runPluginVerifier {
@@ -581,13 +592,16 @@ tasks {
 
   named<Copy>("processIntegrationTestResources") {
     from(sourceSets["integrationTest"].resources)
-    into("$buildDir/resources/integrationTest")
+    into("${layout.buildDirectory.asFile.get()}/resources/integrationTest")
     exclude("**/.idea/**")
     exclude("**/*.xml")
     duplicatesStrategy = DuplicatesStrategy.EXCLUDE
   }
 
-  withType<Test> { systemProperty("idea.test.src.dir", "$buildDir/resources/integrationTest") }
+  withType<Test> {
+    systemProperty(
+        "idea.test.src.dir", "${layout.buildDirectory.asFile.get()}/resources/integrationTest")
+  }
 
   named("classpathIndexCleanup") { dependsOn("processIntegrationTestResources") }
 

diff --git a/src/main/java/com/sourcegraph/cody/agent/protocol/TextDocumentShowOptions.kt b/src/main/java/com/sourcegraph/cody/agent/protocol/TextDocumentShowOptions.kt
@@ -1,5 +1,7 @@
 package com.sourcegraph.cody.agent.protocol
 
+import com.sourcegraph.cody.agent.protocol_generated.Range
+
 data class TextDocumentShowOptions(
     val preserveFocus: Boolean?,
     val preview: Boolean?,

diff --git a/src/main/java/com/sourcegraph/cody/vscode/IntelliJTextDocument.java b/src/main/java/com/sourcegraph/cody/vscode/IntelliJTextDocument.java
@@ -8,8 +8,9 @@
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.util.TextRange;
 import com.intellij.openapi.vfs.VirtualFile;
-import com.sourcegraph.cody.agent.protocol.Position;
-import com.sourcegraph.cody.agent.protocol.Range;
+import com.sourcegraph.cody.agent.protocol_extensions.PositionKt;
+import com.sourcegraph.cody.agent.protocol_generated.Position;
+import com.sourcegraph.cody.agent.protocol_generated.Range;
 import java.net.URI;
 import java.util.Optional;
 import org.jetbrains.annotations.NotNull;
@@ -40,7 +41,7 @@ public String fileName() {
 
   @Override
   public int offsetAt(Position position) {
-    return position.toOffset(this.editor.getDocument());
+    return PositionKt.toOffset(position, this.editor.getDocument());
   }
 
   @Override
@@ -52,7 +53,7 @@ public String getText() {
   public String getText(Range range) {
     return this.editor
         .getDocument()
-        .getText(TextRange.create(offsetAt(range.start), offsetAt(range.end)));
+        .getText(TextRange.create(offsetAt(range.getStart()), offsetAt(range.getEnd())));
   }
 
   @Override

diff --git a/src/main/java/com/sourcegraph/cody/vscode/TextDocument.java b/src/main/java/com/sourcegraph/cody/vscode/TextDocument.java
@@ -1,7 +1,7 @@
 package com.sourcegraph.cody.vscode;
 
-import com.sourcegraph.cody.agent.protocol.Position;
-import com.sourcegraph.cody.agent.protocol.Range;
+import com.sourcegraph.cody.agent.protocol_generated.Position;
+import com.sourcegraph.cody.agent.protocol_generated.Range;
 import java.net.URI;
 import java.util.Optional;
 import org.jetbrains.annotations.NotNull;

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
@@ -10,6 +10,9 @@ import com.intellij.openapi.util.SystemInfoRt
 import com.intellij.util.net.HttpConfigurable
 import com.intellij.util.system.CpuArch
 import com.sourcegraph.cody.agent.protocol.*
+import com.sourcegraph.cody.agent.protocol_generated.ClientCapabilities
+import com.sourcegraph.cody.agent.protocol_generated.ClientInfo
+import com.sourcegraph.cody.agent.protocol_generated.ProtocolTypeAdapters
 import com.sourcegraph.cody.vscode.CancellationToken
 import com.sourcegraph.config.ConfigUtil
 import java.io.*
@@ -107,19 +110,22 @@ private constructor(
           return server
               .initialize(
                   ClientInfo(
+                      name = "JetBrains",
                       version = ConfigUtil.getPluginVersion(),
                       ideVersion = ApplicationInfo.getInstance().build.toString(),
                       workspaceRootUri =
                           ConfigUtil.getWorkspaceRootPath(project).toUri().toString(),
                       extensionConfiguration = ConfigUtil.getAgentConfiguration(project),
                       capabilities =
                           ClientCapabilities(
-                              edit = "enabled",
-                              editWorkspace = "enabled",
-                              codeLenses = "enabled",
-                              showDocument = "enabled",
-                              ignore = "enabled",
-                              untitledDocuments = "enabled")))
+                              edit = ClientCapabilities.EditEnum.Enabled,
+                              editWorkspace = ClientCapabilities.EditWorkspaceEnum.Enabled,
+                              codeLenses = ClientCapabilities.CodeLensesEnum.Enabled,
+                              showDocument = ClientCapabilities.ShowDocumentEnum.Enabled,
+                              ignore = ClientCapabilities.IgnoreEnum.Enabled,
+                              untitledDocuments = ClientCapabilities.UntitledDocumentsEnum.Enabled,
+                              codeActions = ClientCapabilities.CodeActionsEnum.Enabled),
+                  ))
               .thenApply { info ->
                 logger.warn("Connected to Cody agent " + info.name)
                 server.initialized()
@@ -265,16 +271,25 @@ private constructor(
     ): Launcher<CodyAgentServer> {
       return Launcher.Builder<CodyAgentServer>()
           .configureGson { gsonBuilder ->
-            gsonBuilder
-                // emit `null` instead of leaving fields undefined because Cody
-                // VSC has many `=== null` checks that return false for undefined fields.
-                .serializeNulls()
-                .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
-                .registerTypeAdapter(ContextItem::class.java, ContextItem.deserializer)
-                .registerTypeAdapter(Speaker::class.java, speakerDeserializer)
-                .registerTypeAdapter(Speaker::class.java, speakerSerializer)
-                .registerTypeAdapter(URI::class.java, uriDeserializer)
-                .registerTypeAdapter(URI::class.java, uriSerializer)
+            run {
+              gsonBuilder
+                  // emit `null` instead of leaving fields undefined because Cody
+                  // VSC has many `=== null` checks that return false for undefined fields.
+                  .serializeNulls()
+                  .registerTypeAdapter(ContextItem::class.java, ContextItem.deserializer)
+                  .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
+                  // TODO: Once all protocols have migrated we can remove these legacy enum
+                  // conversions
+                  .registerTypeAdapter(Speaker::class.java, speakerDeserializer)
+                  .registerTypeAdapter(Speaker::class.java, speakerSerializer)
+                  .registerTypeAdapter(URI::class.java, uriDeserializer)
+                  .registerTypeAdapter(URI::class.java, uriSerializer)
+
+              ProtocolTypeAdapters.register(gsonBuilder)
+              // This ensures that by default all enums are always serialized to their string
+              // equivalents
+              gsonBuilder.registerTypeAdapterFactory(EnumTypeAdapterFactory())
+            }
           }
           .setRemoteInterface(CodyAgentServer::class.java)
           .traceMessages(traceWriter())

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentCodebase.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentCodebase.kt
@@ -31,7 +31,7 @@ class CodyAgentCodebase(val project: Project) {
       if (repositoryName != null && inferredUrl.getNow(null) != repositoryName) {
         inferredUrl.complete(repositoryName)
         CodyAgentService.withAgent(project) {
-          it.server.configurationDidChange(ConfigUtil.getAgentConfiguration(project))
+          it.server.extensionConfiguration_didChange(ConfigUtil.getAgentConfiguration(project))
         }
       }
     }

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
@@ -1,4 +1,4 @@
-@file:Suppress("FunctionName")
+@file:Suppress("FunctionName", "REDUNDANT_NULLABLE")
 
 package com.sourcegraph.cody.agent
 
@@ -11,13 +11,10 @@ import com.sourcegraph.cody.agent.protocol.ChatModelsParams
 import com.sourcegraph.cody.agent.protocol.ChatModelsResponse
 import com.sourcegraph.cody.agent.protocol.ChatRestoreParams
 import com.sourcegraph.cody.agent.protocol.ChatSubmitMessageParams
-import com.sourcegraph.cody.agent.protocol.ClientInfo
 import com.sourcegraph.cody.agent.protocol.CompletionItemParams
 import com.sourcegraph.cody.agent.protocol.CurrentUserCodySubscription
 import com.sourcegraph.cody.agent.protocol.Event
 import com.sourcegraph.cody.agent.protocol.GetFeatureFlag
-import com.sourcegraph.cody.agent.protocol.GetRepoIdsParam
-import com.sourcegraph.cody.agent.protocol.GetRepoIdsResponse
 import com.sourcegraph.cody.agent.protocol.IgnorePolicySpec
 import com.sourcegraph.cody.agent.protocol.IgnoreTestParams
 import com.sourcegraph.cody.agent.protocol.IgnoreTestResponse
@@ -28,14 +25,8 @@ import com.sourcegraph.cody.agent.protocol.RemoteRepoHasParams
 import com.sourcegraph.cody.agent.protocol.RemoteRepoHasResponse
 import com.sourcegraph.cody.agent.protocol.RemoteRepoListParams
 import com.sourcegraph.cody.agent.protocol.RemoteRepoListResponse
-import com.sourcegraph.cody.agent.protocol.ServerInfo
 import com.sourcegraph.cody.agent.protocol.TelemetryEvent
-import com.sourcegraph.cody.agent.protocol_generated.EditTask
-import com.sourcegraph.cody.agent.protocol_generated.EditTask_AcceptParams
-import com.sourcegraph.cody.agent.protocol_generated.EditTask_CancelParams
-import com.sourcegraph.cody.agent.protocol_generated.EditTask_GetTaskDetailsParams
-import com.sourcegraph.cody.agent.protocol_generated.EditTask_RetryParams
-import com.sourcegraph.cody.agent.protocol_generated.EditTask_UndoParams
+import com.sourcegraph.cody.agent.protocol_generated.*
 import com.sourcegraph.cody.chat.ConnectionId
 import java.util.concurrent.CompletableFuture
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
@@ -50,15 +41,36 @@ interface _SubsetGeneratedCodyAgentServer {
   // ========
   // Requests
   // ========
+  @JsonRequest("initialize") fun initialize(params: ClientInfo): CompletableFuture<ServerInfo>
+
   @JsonRequest("editTask/retry")
   fun editTask_retry(params: EditTask_RetryParams): CompletableFuture<EditTask>
 
   @JsonRequest("editTask/getTaskDetails")
   fun editTask_getTaskDetails(params: EditTask_GetTaskDetailsParams): CompletableFuture<EditTask>
 
-  // =============
-  // Notifications
-  // =============
+  @JsonRequest("diagnostics/publish")
+  fun diagnostics_publish(params: Diagnostics_PublishParams): CompletableFuture<Null?>
+
+  @JsonRequest("codeActions/provide")
+  fun codeActions_provide(
+      params: CodeActions_ProvideParams
+  ): CompletableFuture<CodeActions_ProvideResult>
+
+  @JsonRequest("codeActions/trigger")
+  fun codeActions_trigger(params: CodeActions_TriggerParams): CompletableFuture<EditTask>
+
+  @JsonRequest("graphql/getRepoIds")
+  fun graphql_getRepoIds(
+      params: Graphql_GetRepoIdsParams
+  ): CompletableFuture<Graphql_GetRepoIdsResult>
+
+  //  // =============
+  //  // Notifications
+  //  // =============
+
+  @JsonNotification("extensionConfiguration/didChange")
+  fun extensionConfiguration_didChange(params: ExtensionConfiguration)
 }
 
 // TODO: Requests waiting to be migrated & tested for compatibility. Avoid placing new protocol
@@ -70,9 +82,6 @@ interface _SubsetGeneratedCodyAgentServer {
  * works similar to JavaScript Proxy.
  */
 interface _LegacyAgentServer {
-  // Requests
-  @JsonRequest("initialize") fun initialize(clientInfo: ClientInfo): CompletableFuture<ServerInfo>
-
   @JsonRequest("shutdown") fun shutdown(): CompletableFuture<Void?>
 
   @JsonRequest("autocomplete/execute")
@@ -85,23 +94,20 @@ interface _LegacyAgentServer {
 
   @JsonRequest("graphql/currentUserId") fun currentUserId(): CompletableFuture<String>
 
-  @JsonRequest("graphql/getRepoIds")
-  fun getRepoIds(repoName: GetRepoIdsParam): CompletableFuture<GetRepoIdsResponse>
-
+  // TODO(CODY-2826): Would be nice if we can generate some set of "known" feature flags from the
+  // protocol
   @JsonRequest("featureFlags/getFeatureFlag")
   fun evaluateFeatureFlag(flagName: GetFeatureFlag): CompletableFuture<Boolean?>
 
+  // TODO(CODY-2827): To avoid having to pass annoying null values we should generate a default
+  // value
   @JsonRequest("graphql/getCurrentUserCodySubscription")
   fun getCurrentUserCodySubscription(): CompletableFuture<CurrentUserCodySubscription?>
 
-  // Notifications
   @JsonNotification("initialized") fun initialized()
 
   @JsonNotification("exit") fun exit()
 
-  @JsonNotification("extensionConfiguration/didChange")
-  fun configurationDidChange(document: ExtensionConfiguration)
-
   @JsonNotification("textDocument/didFocus")
   fun textDocumentDidFocus(document: ProtocolTextDocument)
 

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/ExtensionConfiguration.kt b/src/main/kotlin/com/sourcegraph/cody/agent/ExtensionConfiguration.kt
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/intellij_extensions/Document.kt b/src/main/kotlin/com/sourcegraph/cody/agent/intellij_extensions/Document.kt
@@ -0,0 +1,30 @@
+package com.sourcegraph.cody.agent.intellij_extensions
+
+import com.intellij.openapi.editor.Document
+import com.sourcegraph.cody.agent.protocol_extensions.Position
+import com.sourcegraph.cody.agent.protocol_generated.Position
+import com.sourcegraph.cody.agent.protocol_generated.Range
+
+fun Document.codyPosition(offset: Int): Position {
+  val line = this.getLineNumber(offset)
+  val lineStartOffset = this.getLineStartOffset(line)
+  val character = offset - lineStartOffset
+  return Position(line, character)
+}
+
+fun Document.codyRange(startOffset: Int, endOffset: Int): Range {
+  val startLine = this.getLineNumber(startOffset)
+  val lineStartOffset1 = this.getLineStartOffset(startLine)
+  val startCharacter = startOffset - lineStartOffset1
+
+  val endLine = this.getLineNumber(endOffset)
+  val lineStartOffset2 =
+      if (startLine == endLine) {
+        lineStartOffset1
+      } else {
+        this.getLineStartOffset(endLine)
+      }
+  val endCharacter = endOffset - lineStartOffset2
+
+  return Range(Position(startLine, startCharacter), Position(endLine, endCharacter))
+}
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AutocompleteParams.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AutocompleteParams.kt
@@ -1,5 +1,7 @@
 package com.sourcegraph.cody.agent.protocol
 
+import com.sourcegraph.cody.agent.protocol_generated.Position
+
 enum class AutocompleteTriggerKind(val value: String) {
   AUTOMATIC("Automatic"),
   INVOKE("Invoke"),