From dd8fe46c1feb161d27dbeb007557dc1b40a2e0d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Kondratek?= <mik.kondratek@gmail.com>
Date: Mon, 25 Nov 2024 11:49:15 +0100
Subject: [PATCH] Use `editCommands/document` instead of custom command

---
 .../kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt   | 3 +++
 .../sourcegraph/cody/edit/actions/DocumentCodeAction.kt    | 7 ++-----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
index 9220da56e5..00da34775b 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
@@ -117,6 +117,9 @@ interface _SubsetGeneratedCodyAgentServer {
   @JsonRequest("editCommands/code")
   fun editCommands_code(params: EditCommands_CodeParams): CompletableFuture<EditTask>
 
+  @JsonRequest("editCommands/document")
+  fun editCommands_document(params: Null?): CompletableFuture<EditTask>
+
   @JsonRequest("webview/resolveWebviewView")
   fun webview_resolveWebviewView(params: Webview_ResolveWebviewViewParams): CompletableFuture<Null?>
 
diff --git a/src/main/kotlin/com/sourcegraph/cody/edit/actions/DocumentCodeAction.kt b/src/main/kotlin/com/sourcegraph/cody/edit/actions/DocumentCodeAction.kt
index 4688d0f3a5..2d2e2b5c3a 100644
--- a/src/main/kotlin/com/sourcegraph/cody/edit/actions/DocumentCodeAction.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/edit/actions/DocumentCodeAction.kt
@@ -1,16 +1,13 @@
 package com.sourcegraph.cody.edit.actions
 
 import com.sourcegraph.cody.agent.CodyAgentService
-import com.sourcegraph.cody.agent.protocol_generated.Commands_CustomParams
-import com.sourcegraph.cody.agent.protocol_generated.CustomEditCommandResult
 
 class DocumentCodeAction :
     BaseEditCodeAction({ editor ->
       editor.project?.let { project ->
         CodyAgentService.withAgent(project) { agent ->
-          val customCommandResult = agent.server.commands_custom(Commands_CustomParams("doc")).get()
-          val result = customCommandResult as? CustomEditCommandResult ?: return@withAgent
-          EditCodeAction.completedEditTasks[result.editResult.id] = result.editResult
+          val result = agent.server.editCommands_document(null).get()
+          EditCodeAction.completedEditTasks[result.id] = result
         }
       }
     }) {