diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
index 1150d04547..ae55481fff 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
@@ -6,13 +6,13 @@ import com.sourcegraph.cody.agent.protocol.DebugMessage
 import com.sourcegraph.cody.agent.protocol.OpenExternalParams
 import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument
 import com.sourcegraph.cody.agent.protocol.RemoteRepoFetchState
-import com.sourcegraph.cody.agent.protocol.TextDocumentEditParams
 import com.sourcegraph.cody.agent.protocol.TextDocumentShowParams
 import com.sourcegraph.cody.agent.protocol.UntitledTextDocument
 import com.sourcegraph.cody.agent.protocol.WebviewCreateWebviewPanelParams
-import com.sourcegraph.cody.agent.protocol.WorkspaceEditParams
 import com.sourcegraph.cody.agent.protocol_generated.DisplayCodeLensParams
 import com.sourcegraph.cody.agent.protocol_generated.EditTask
+import com.sourcegraph.cody.agent.protocol_generated.TextDocumentEditParams
+import com.sourcegraph.cody.agent.protocol_generated.WorkspaceEditParams
 import com.sourcegraph.cody.ui.NativeWebviewProvider
 import java.util.concurrent.CompletableFuture
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextDocumentEditParams.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextDocumentEditParams.kt
deleted file mode 100644
index ee57c1d103..0000000000
--- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextDocumentEditParams.kt
+++ /dev/null
@@ -1,7 +0,0 @@
-package com.sourcegraph.cody.agent.protocol
-
-data class TextDocumentEditParams(
-    val uri: String,
-    val edits: List<TextEdit>,
-    val options: TextDocumentEditOptions? = null
-)
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextEdit.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextEdit.kt
deleted file mode 100644
index e01906db9f..0000000000
--- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/TextEdit.kt
+++ /dev/null
@@ -1,18 +0,0 @@
-package com.sourcegraph.cody.agent.protocol
-
-import com.sourcegraph.cody.agent.protocol_generated.Position
-import com.sourcegraph.cody.agent.protocol_generated.Range
-
-data class TextEdit(
-    // This tag will be 'replace', 'insert', or 'delete'.
-    val type: String,
-
-    // Valid for replace & delete.
-    val range: Range? = null,
-
-    // Valid for insert.
-    val position: Position? = null,
-
-    // Valid for replace & insert.
-    val value: String? = null
-)
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditOperation.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditOperation.kt
deleted file mode 100644
index bc082f10dc..0000000000
--- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditOperation.kt
+++ /dev/null
@@ -1,18 +0,0 @@
-package com.sourcegraph.cody.agent.protocol
-
-data class WorkspaceEditParamsOptions(
-    val overwrite: Boolean = false,
-    val ignoreIfNotExists: Boolean = false,
-    val recursive: Boolean = false
-)
-
-data class WorkspaceEditOperation(
-    val type: String, // all
-    val uri: String? = null, // created, delete, edit
-    val oldUri: String? = null, // rename
-    val newUri: String? = null, // rename
-    val textContents: String? = null, // create-file
-    val options: WorkspaceEditParamsOptions? = null, // all
-    val metadata: WorkspaceEditMetadata? = null, // all
-    val edits: List<TextEdit>? = null
-) // edit
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditParams.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditParams.kt
deleted file mode 100644
index 41675d7dba..0000000000
--- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/WorkspaceEditParams.kt
+++ /dev/null
@@ -1,6 +0,0 @@
-package com.sourcegraph.cody.agent.protocol
-
-data class WorkspaceEditParams(
-    val operations: List<WorkspaceEditOperation>,
-    val metadata: WorkspaceEditMetadata? = null
-)
diff --git a/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt b/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
index 4a5d9e49f6..b1337506b5 100644
--- a/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/edit/EditService.kt
@@ -6,9 +6,16 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.project.Project
-import com.sourcegraph.cody.agent.protocol.TextEdit
-import com.sourcegraph.cody.agent.protocol.WorkspaceEditParams
 import com.sourcegraph.cody.agent.protocol_extensions.toOffset
+import com.sourcegraph.cody.agent.protocol_generated.CreateFileOperation
+import com.sourcegraph.cody.agent.protocol_generated.DeleteFileOperation
+import com.sourcegraph.cody.agent.protocol_generated.DeleteTextEdit
+import com.sourcegraph.cody.agent.protocol_generated.EditFileOperation
+import com.sourcegraph.cody.agent.protocol_generated.InsertTextEdit
+import com.sourcegraph.cody.agent.protocol_generated.RenameFileOperation
+import com.sourcegraph.cody.agent.protocol_generated.ReplaceTextEdit
+import com.sourcegraph.cody.agent.protocol_generated.TextEdit
+import com.sourcegraph.cody.agent.protocol_generated.WorkspaceEditParams
 import com.sourcegraph.utils.CodyEditorUtil
 
 @Service(Service.Level.PROJECT)
@@ -38,31 +45,23 @@ class EditService(val project: Project) {
 
     return WriteCommandAction.runWriteCommandAction<Boolean>(project) {
       edits.reversed().all { edit ->
-        when (edit.type) {
-          "replace",
-          "delete" -> {
-            if (edit.range != null) {
-              document.replaceString(
-                  edit.range.start.toOffset(document),
-                  edit.range.end.toOffset(document),
-                  edit.value ?: "")
-              true
-            } else {
-              logger.warn("Edit range is null for ${edit.type} operation")
-              false
-            }
+        when (edit) {
+          is ReplaceTextEdit -> {
+            document.replaceString(
+                edit.range.start.toOffset(document), edit.range.end.toOffset(document), edit.value)
+            true
           }
-          "insert" -> {
-            if (edit.position != null) {
-              document.insertString(edit.position.toOffset(document), edit.value ?: "")
-              true
-            } else {
-              logger.warn("Edit position is null for insert operation")
-              false
-            }
+          is DeleteTextEdit -> {
+            document.deleteString(
+                edit.range.start.toOffset(document), edit.range.end.toOffset(document))
+            true
+          }
+          is InsertTextEdit -> {
+            document.insertString(edit.position.toOffset(document), edit.value)
+            true
           }
           else -> {
-            logger.warn("Unknown edit type: ${edit.type}")
+            logger.warn("Unknown edit type: $edit")
             false
           }
         }
@@ -73,34 +72,25 @@ class EditService(val project: Project) {
   fun performWorkspaceEdit(workspaceEditParams: WorkspaceEditParams): Boolean {
     return workspaceEditParams.operations.all { op ->
       // TODO: We need to support the file-level operations.
-      when (op.type) {
-        "create-file" -> {
+      when (op) {
+        is CreateFileOperation -> {
           logger.warn("Workspace edit operation created a file: ${op.uri}")
           return false
         }
-        "rename-file" -> {
+        is RenameFileOperation -> {
           logger.warn("Workspace edit operation renamed a file: ${op.oldUri} -> ${op.newUri}")
           return false
         }
-        "delete-file" -> {
+        is DeleteFileOperation -> {
           logger.warn("Workspace edit operation deleted a file: ${op.uri}")
           return false
         }
-        "edit-file" -> {
-          if (op.edits == null) {
-            logger.warn("Workspace edit operation has no edits")
-            return false
-          } else if (op.uri == null) {
-            logger.warn("Workspace edit operation has null uri")
-            return false
-          } else {
-            logger.info("Applying workspace edit to a file: ${op.uri}")
-            performTextEdits(op.uri, op.edits)
-          }
+        is EditFileOperation -> {
+          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.type}")
+          logger.warn("DocumentCommand session received unknown workspace edit operation: $op")
           return false
         }
       }