From 0d251bc6c173fdba35a3af0f2ec2085d4dcf9158 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Kondratek?= <mik.kondratek@gmail.com>
Date: Wed, 14 Aug 2024 19:57:41 +0200
Subject: [PATCH] Add handling for create-file workspace op

---
 .../com/sourcegraph/cody/agent/CodyAgent.kt   |  2 +-
 .../sourcegraph/cody/agent/CodyAgentClient.kt | 26 ++++++++++++++++++-
 .../com/sourcegraph/cody/edit/EditService.kt  | 15 ++++-------
 .../cody/agent/CodyAgentClientTest.kt         |  2 +-
 4 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
index b5af836d82..4937d79331 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
@@ -103,7 +103,7 @@ private constructor(
     fun create(project: Project): CompletableFuture<CodyAgent> {
       try {
         val conn = startAgentProcess()
-        val client = CodyAgentClient(WebUIServiceWebviewProvider(project), project)
+        val client = CodyAgentClient(project, WebUIServiceWebviewProvider(project))
         client.onSetConfigFeatures = project.service<CurrentConfigFeatures>()
         val launcher = startAgentLauncher(conn, client)
         val server = launcher.remoteProxy
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
index ae55481fff..d4f5aeca10 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
@@ -1,7 +1,14 @@
 package com.sourcegraph.cody.agent
 
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.diagnostic.Logger
+import com.intellij.openapi.fileChooser.FileChooserFactory
+import com.intellij.openapi.fileChooser.FileSaverDescriptor
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.project.guessProjectDir
+import com.intellij.openapi.vfs.VfsUtil
+import com.intellij.openapi.vfs.VirtualFile
 import com.sourcegraph.cody.agent.protocol.DebugMessage
 import com.sourcegraph.cody.agent.protocol.OpenExternalParams
 import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
@@ -23,7 +30,7 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
  * and notifications sent by the agent.
  */
 @Suppress("unused")
-class CodyAgentClient(private val webview: NativeWebviewProvider) {
+class CodyAgentClient(private val project: Project, private val webview: NativeWebviewProvider) {
   companion object {
     private val logger = Logger.getInstance(CodyAgentClient::class.java)
   }
@@ -253,4 +260,21 @@ class CodyAgentClient(private val webview: NativeWebviewProvider) {
 
     logger.debug("webview/postMessage ${params.id}: ${params.message}")
   }
+
+  @JsonRequest("window/showSaveDialog")
+  fun window_showSaveDialog(): CompletableFuture<String> {
+    var outputDir: VirtualFile? = project.guessProjectDir()
+    if (outputDir == null || !outputDir.exists()) {
+      outputDir = VfsUtil.getUserHomeDir()
+    }
+
+    val descriptor = FileSaverDescriptor("Cody: Save as New File", "Save file")
+    val saveFileFuture = CompletableFuture<String>()
+    runInEdt {
+      val dialog = FileChooserFactory.getInstance().createSaveFileDialog(descriptor, project)
+      val result = dialog.save(outputDir, "Untitled")
+      saveFileFuture.complete(result?.file?.path)
+    }
+    return saveFileFuture
+  }
 }
diff --git a/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt b/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
index b1337506b5..553fd9e707 100644
--- a/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
@@ -60,10 +60,6 @@ class EditService(val project: Project) {
             document.insertString(edit.position.toOffset(document), edit.value)
             true
           }
-          else -> {
-            logger.warn("Unknown edit type: $edit")
-            false
-          }
         }
       }
     }
@@ -74,8 +70,11 @@ class EditService(val project: Project) {
       // TODO: We need to support the file-level operations.
       when (op) {
         is CreateFileOperation -> {
-          logger.warn("Workspace edit operation created a file: ${op.uri}")
-          return false
+          logger.info("Workspace edit operation created a file: ${op.uri}")
+          val file =
+              CodyEditorUtil.createFileOrScratchFromUntitled(project, op.uri, content = "")
+                  ?: return false
+          CodyEditorUtil.showDocument(project, file)
         }
         is RenameFileOperation -> {
           logger.warn("Workspace edit operation renamed a file: ${op.oldUri} -> ${op.newUri}")
@@ -89,10 +88,6 @@ class EditService(val project: Project) {
           logger.info("Applying workspace edit to a file: ${op.uri}")
           performTextEdits(op.uri, op.edits)
         }
-        else -> {
-          logger.warn("DocumentCommand session received unknown workspace edit operation: $op")
-          return false
-        }
       }
     }
   }
diff --git a/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt b/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt
index f5ec841426..ff7a359c88 100644
--- a/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt
+++ b/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt
@@ -21,7 +21,7 @@ class CodyAgentClientTest : BasePlatformTestCase() {
   private val condition = lock.newCondition()
 
   private fun client(): CodyAgentClient {
-    val client = CodyAgentClient(StubWebviewProvider(), project)
+    val client = CodyAgentClient(project, StubWebviewProvider())
     client.onSetConfigFeatures = ConfigFeaturesObserver {
       lock.lock()
       try {