V1 new client extension API, for inline edits + document code (#3445)

Co-authored-by: Dominic Cooney <dominic.cooney@sourcegraph.com>
Co-authored-by: Beatrix Woo <abeatrix@sourcegraph.com>

.vscode/launch.json

@@ -62,6 +62,13 @@
       "outFiles": ["${workspaceFolder}/vscode/dist/**/*.js"],
       "args": ["--extensionDevelopmentPath=${workspaceRoot}/vscode", "--extensionDevelopmentKind=web"]
     },
+    {
+      "name": "Attach to Agent",
+      "port": 9229,
+      "request": "attach",
+      "skipFiles": ["<node_internals>/**"],
+      "type": "node"
+    },
     {
       "type": "node",
       "request": "launch",

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/CodyAgentClient.kt

@@ -28,8 +28,10 @@ interface CodyAgentClient {
   // =============
   @JsonNotification("debug/message")
   fun debug_message(params: DebugMessage)
-  @JsonNotification("editTaskState/didChange")
-  fun editTaskState_didChange(params: EditTask)
+  @JsonNotification("editTask/didUpdate")
+  fun editTask_didUpdate(params: EditTask)
+  @JsonNotification("editTask/didDelete")
+  fun editTask_didDelete(params: EditTask)
   @JsonNotification("codeLenses/display")
   fun codeLenses_display(params: DisplayCodeLensParams)
   @JsonNotification("webview/postMessage")

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/CodyAgentServer.kt

@@ -34,10 +34,20 @@ interface CodyAgentServer {
   fun commands_smell(params: Null): CompletableFuture<String>
   @JsonRequest("commands/custom")
   fun commands_custom(params: Commands_CustomParams): CompletableFuture<CustomCommandResult>
+  @JsonRequest("editCommands/code")
+  fun editCommands_code(params: EditCommands_CodeParams): CompletableFuture<EditTask>
   @JsonRequest("editCommands/test")
   fun editCommands_test(params: Null): CompletableFuture<EditTask>
   @JsonRequest("commands/document")
   fun commands_document(params: Null): CompletableFuture<EditTask>
+  @JsonRequest("editTask/accept")
+  fun editTask_accept(params: FixupTaskID): CompletableFuture<Null>
+  @JsonRequest("editTask/undo")
+  fun editTask_undo(params: FixupTaskID): CompletableFuture<Null>
+  @JsonRequest("editTask/cancel")
+  fun editTask_cancel(params: FixupTaskID): CompletableFuture<Null>
+  @JsonRequest("editTask/getFoldingRanges")
+  fun editTask_getFoldingRanges(params: GetFoldingRangeParams): CompletableFuture<GetFoldingRangeResult>
   @JsonRequest("command/execute")
   fun command_execute(params: ExecuteCommandParams): CompletableFuture<Any>
   @JsonRequest("autocomplete/execute")

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/EditCommands_CodeParams.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.protocol_generated
+
+data class EditCommands_CodeParams(
+  val params: ParamsParams? = null,
+)
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/EditTask.kt

@@ -5,5 +5,6 @@ data class EditTask(
   val id: String? = null,
   val state: CodyTaskState? = null,
   val error: CodyError? = null,
+  val selectionRange: Range? = null,
 )
 

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/FixupTaskID.kt

@@ -0,0 +1,5 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.protocol_generated
+
+typealias FixupTaskID = String // One of: 
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/GetFoldingRangeParams.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.protocol_generated
+
+data class GetFoldingRangeParams(
+  val uri: String? = null,
+)
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/GetFoldingRangeResult.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.protocol_generated
+
+data class GetFoldingRangeResult(
+  val ranges: List<Range>? = null,
+)
+

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/ParamsParams.kt

@@ -0,0 +1,7 @@
+@file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
+package com.sourcegraph.cody.protocol_generated
+
+data class ParamsParams(
+  val instruction: String? = null,
+)
+

agent/package.json

@@ -17,7 +17,7 @@
     "build": "pnpm run build:root && pnpm run build:agent",
     "build-minify": "pnpm run build --minify",
     "agent": "pnpm run build && node dist/index.js",
-    "agent:debug": "pnpm run build && node --inspect --enable-source-maps ./dist/index.js",
+    "agent:debug": "pnpm run build && CODY_AGENT_TRACE_PATH=/tmp/agent.json CODY_AGENT_DEBUG_REMOTE=true node --enable-source-maps ./dist/index.js",
     "build-ts": "tsc --build",
     "build-agent-binaries": "pnpm run build && pkg  --compress GZip --no-bytecode --public-packages \"*\" --public . --out-path ${AGENT_EXECUTABLE_TARGET_DIRECTORY:-dist}",
     "test-agent-binary": "esbuild ./scripts/test-agent-binary.ts --bundle --platform=node --alias:vscode=./src/vscode-shim.ts --outdir=dist && node ./dist/test-agent-binary.js",

agent/src/AgentFixupControls.ts

@@ -0,0 +1,56 @@
+import type { FixupFile } from '../../vscode/src/non-stop/FixupFile'
+import type { FixupTask, FixupTaskID } from '../../vscode/src/non-stop/FixupTask'
+import type { FixupActor, FixupFileCollection } from '../../vscode/src/non-stop/roles'
+import type { FixupControlApplicator } from '../../vscode/src/non-stop/strategies'
+import { type Agent, errorToCodyError } from './agent'
+import type { EditTask } from './protocol-alias'
+
+export class AgentFixupControls implements FixupControlApplicator {
+    constructor(
+        private readonly fixups: FixupActor & FixupFileCollection,
+        private readonly notify: typeof Agent.prototype.notify
+    ) {}
+
+    public accept(id: FixupTaskID): void {
+        const task = this.fixups.taskForId(id)
+        if (task) {
+            this.fixups.accept(task)
+        }
+    }
+
+    public undo(id: FixupTaskID): void {
+        const task = this.fixups.taskForId(id)
+        if (task) {
+            this.fixups.undo(task)
+        }
+    }
+
+    public cancel(id: FixupTaskID): void {
+        const task = this.fixups.taskForId(id)
+        if (task) {
+            this.fixups.cancel(task)
+        }
+    }
+
+    // FixupControlApplicator
+
+    didUpdateTask(task: FixupTask): void {
+        this.notify('editTask/didUpdate', AgentFixupControls.serialize(task))
+    }
+    didDeleteTask(task: FixupTask): void {
+        this.notify('editTask/didDelete', AgentFixupControls.serialize(task))
+    }
+
+    visibleFilesWithTasksMaybeChanged(files: readonly FixupFile[]): void {}
+
+    dispose() {}
+
+    public static serialize(task: FixupTask): EditTask {
+        return {
+            id: task.id,
+            state: task.state,
+            error: errorToCodyError(task.error),
+            selectionRange: task.selectionRange,
+        }
+    }
+}

agent/src/TestClient.ts

@@ -343,33 +343,52 @@ export class TestClient extends MessageHandler {
                 )
         }
 
-        let disposable: vscode.Disposable
+        let disposables: vscode.Disposable[]
         return new Promise<void>((resolve, reject) => {
-            disposable = this.onDidChangeTaskState(({ id, state, error }) => {
-                if (id === params.id) {
-                    switch (state) {
-                        case CodyTaskState.applied:
-                            return resolve()
-                        case CodyTaskState.error:
-                        case CodyTaskState.finished:
-                            return reject(
-                                new Error(
-                                    `Task reached terminal state before being applied ${JSON.stringify({
-                                        id,
-                                        state: CodyTaskState[state],
-                                        error,
-                                    })}`
+            disposables = [
+                this.onDidUpdateTask(({ id, state, error }) => {
+                    if (id === params.id) {
+                        switch (state) {
+                            case CodyTaskState.applied:
+                                return resolve()
+                            case CodyTaskState.error:
+                            case CodyTaskState.finished:
+                                return reject(
+                                    new Error(
+                                        `Task reached terminal state before being applied ${JSON.stringify(
+                                            {
+                                                id,
+                                                state: CodyTaskState[state],
+                                                error,
+                                            }
+                                        )}`
+                                    )
                                 )
-                            )
+                        }
                     }
-                }
-            })
-        }).finally(() => disposable.dispose())
+                }),
+                this.onDidDeleteTask(task => {
+                    if (task.id === params.id) {
+                        // Applied tasks can also be deleted, but in that case
+                        // the Promise is already resolved and this is a no-op.
+                        reject(
+                            new Error(`Task was deleted before being applied ${JSON.stringify(task)}`)
+                        )
+                    }
+                }),
+            ]
+        }).finally(() => {
+            for (const disposable of disposables) {
+                disposable.dispose()
+            }
+        })
     }
 
     public codeLenses = new Map<string, ProtocolCodeLens[]>()
-    public newTaskState = new vscode.EventEmitter<EditTask>()
-    public onDidChangeTaskState = this.newTaskState.event
+    public taskUpdate = new vscode.EventEmitter<EditTask>()
+    public onDidUpdateTask = this.taskUpdate.event
+    public taskDelete = new vscode.EventEmitter<EditTask>()
+    public onDidDeleteTask = this.taskDelete.event
     public webviewMessages: WebviewPostMessageParams[] = []
     public webviewMessagesEmitter = new vscode.EventEmitter<WebviewPostMessageParams>()
 
@@ -416,8 +435,11 @@ export class TestClient extends MessageHandler {
         this.connectProcess(this.agentProcess, error => {
             console.error(error)
         })
-        this.registerNotification('editTaskState/didChange', params => {
-            this.newTaskState.fire(params)
+        this.registerNotification('editTask/didUpdate', params => {
+            this.taskUpdate.fire(params)
+        })
+        this.registerNotification('editTask/didDelete', params => {
+            this.taskDelete.fire(params)
         })
 
         this.registerNotification('webview/postMessage', params => {