Agent: add support for client-side progress bars (#2658)

Fixes sourcegraph/jetbrains#257

Previously, the agent didn't support rendering progress bars on the
client. This was problematic because symf symbol uses progress bars to
inform the user when symf is indexing, and allows users to cancel
indexing. This PR closes this feature gap by adding four new JSON-RPC
notification endpoints to the agent protocol to deal with progress bars:

* `progress/start`: server to client.
* `progress/report`: server to client.
* `progress/end`: server to client.
* `progress/end`: client to server. Only valid for cancelable progress
bars.

Agent clients must add `progressBars: 'enabled'` to the client
capabilities during the `initalize` handshake to receive
progress-related notifications.

Additionally, to help clients with testing progress bars, there are two
new testing-related JSON-RPC endpoints

* `testing/progress`. Client request to server.
* `testing/progressCancelation`. Client request to server.

This PR adds two test cases that demonstrate how progress bars work,
including how to cancel a progress bar.

## Test plan

See two new test cases for progress bars.

<!-- Required. See
https://sourcegraph.com/docs/dev/background-information/testing_principles.
-->

agent/src/agent.ts

@@ -281,6 +281,59 @@ export class Agent extends MessageHandler {
             })
         })
 
+        this.registerNotification('progress/cancel', ({ id }) => {
+            const token = vscode_shim.progressBars.get(id)
+            if (token) {
+                token.cancel()
+            } else {
+                console.error(`progress/cancel: unknown ID ${id}`)
+            }
+        })
+
+        this.registerRequest('testing/progress', async ({ title }) => {
+            const thenable = await vscode.window.withProgress(
+                { title: 'testing/progress', location: vscode.ProgressLocation.Notification, cancellable: true },
+                progress => {
+                    progress.report({ message: 'message1' })
+                    progress.report({ increment: 50 })
+                    progress.report({ increment: 50 })
+                    return Promise.resolve({ result: `Hello ${title}` })
+                }
+            )
+            return thenable
+        })
+
+        this.registerRequest('testing/progressCancelation', async ({ title }) => {
+            const message = await vscode.window.withProgress<string>(
+                {
+                    title: 'testing/progressCancelation',
+                    location: vscode.ProgressLocation.Notification,
+                    cancellable: true,
+                },
+                (progress, token) => {
+                    return new Promise<string>((resolve, reject) => {
+                        token.onCancellationRequested(() => {
+                            progress.report({ message: 'before resolution' })
+                            resolve(`request with title '${title}' cancelled`)
+                            progress.report({ message: 'after resolution' })
+                        })
+                        setTimeout(
+                            () =>
+                                reject(
+                                    new Error(
+                                        'testing/progressCancelation did not resolve within 5 seconds. ' +
+                                            'To fix this problem, send a progress/cancel notification with the same ID ' +
+                                            'as the progress/start notification with title "testing/progressCancelation"'
+                                    )
+                                ),
+                            5_000
+                        )
+                    })
+                }
+            )
+            return { result: message }
+        })
+
         this.registerRequest('recipes/list', () =>
             Promise.resolve(
                 Object.values<RecipeInfo>(registeredRecipes).map(({ id, title }) => ({

agent/src/index.test.ts

@@ -5,6 +5,7 @@ import os from 'os'
 import path from 'path'
 
 import { afterAll, beforeAll, describe, expect, it } from 'vitest'
+import * as vscode from 'vscode'
 import { Uri } from 'vscode'
 
 import { type ChatMessage, type ContextFile } from '@sourcegraph/cody-shared'
@@ -13,11 +14,33 @@ import type { ExtensionMessage } from '../../vscode/src/chat/protocol'
 
 import { AgentTextDocument } from './AgentTextDocument'
 import { MessageHandler } from './jsonrpc-alias'
-import { type ClientInfo, type ServerInfo } from './protocol-alias'
+import { type ClientInfo, type ProgressReportParams, type ProgressStartParams, type ServerInfo } from './protocol-alias'
+
+type ProgressMessage = ProgressStartMessage | ProgressReportMessage | ProgressEndMessage
+interface ProgressStartMessage {
+    method: 'progress/start'
+    id: string
+    message: ProgressStartParams
+}
+interface ProgressReportMessage {
+    method: 'progress/report'
+    id: string
+    message: ProgressReportParams
+}
+interface ProgressEndMessage {
+    method: 'progress/end'
+    id: string
+    message: {}
+}
 
 export class TestClient extends MessageHandler {
     public info: ClientInfo
     public agentProcess?: ChildProcessWithoutNullStreams
+    // Array of all raw `progress/*` notification. Typed as `any` because
+    // start/end/report have different types.
+    public progressMessages: ProgressMessage[] = []
+    public progressIDs = new Map<string, number>()
+    public progressStartEvents = new vscode.EventEmitter<ProgressStartParams>()
 
     constructor(
         public readonly name: string,
@@ -28,12 +51,34 @@ export class TestClient extends MessageHandler {
         this.name = name
         this.info = this.getClientInfo()
 
+        this.registerNotification('progress/start', message => {
+            this.progressStartEvents.fire(message)
+            message.id = this.progressID(message.id)
+            this.progressMessages.push({ method: 'progress/start', id: message.id, message })
+        })
+        this.registerNotification('progress/report', message => {
+            message.id = this.progressID(message.id)
+            this.progressMessages.push({ method: 'progress/report', id: message.id, message })
+        })
+        this.registerNotification('progress/end', ({ id }) => {
+            this.progressMessages.push({ method: 'progress/end', id: this.progressID(id), message: {} })
+        })
         this.registerNotification('debug/message', message => {
             // Uncomment below to see `logDebug` messages.
             // console.log(`${message.channel}: ${message.message}`)
         })
     }
 
+    private progressID(id: string): string {
+        const fromCache = this.progressIDs.get(id)
+        if (fromCache !== undefined) {
+            return `ID_${fromCache}`
+        }
+        const freshID = this.progressIDs.size
+        this.progressIDs.set(id, freshID)
+        return `ID_${freshID}`
+    }
+
     public async initialize() {
         this.agentProcess = this.spawnAgentProcess()
 
@@ -148,6 +193,9 @@ export class TestClient extends MessageHandler {
             version: 'v1',
             workspaceRootUri: workspaceRootUri.toString(),
             workspaceRootPath: workspaceRootUri.fsPath,
+            capabilities: {
+                progressBars: 'enabled',
+            },
             extensionConfiguration: {
                 anonymousUserID: this.name + 'abcde1234',
                 accessToken: this.accessToken ?? 'sgp_RRRRRRRREEEEEEEDDDDDDAAACCCCCTEEEEEEEDDD',
@@ -406,6 +454,79 @@ describe('Agent', () => {
         await client.request('recipes/execute', { id: 'chat-question', humanChatInput: 'How do I implement sum?' })
     }, 600)
 
+    describe('progress bars', () => {
+        it('messages are sent', async () => {
+            const { result } = await client.request('testing/progress', { title: 'Susan' })
+            expect(result).toStrictEqual('Hello Susan')
+            let progressID: string | undefined
+            for (const message of client.progressMessages) {
+                if (message.method === 'progress/start' && message.message.options.title === 'testing/progress') {
+                    progressID = message.message.id
+                    break
+                }
+            }
+            assert(progressID !== undefined, JSON.stringify(client.progressMessages))
+            const messages = client.progressMessages
+                .filter(message => message.id === progressID)
+                .map(({ method, message }) => [method, { ...message, id: 'THE_ID' }])
+            expect(messages).toMatchInlineSnapshot(`
+              [
+                [
+                  "progress/start",
+                  {
+                    "id": "THE_ID",
+                    "options": {
+                      "cancellable": true,
+                      "location": "Notification",
+                      "title": "testing/progress",
+                    },
+                  },
+                ],
+                [
+                  "progress/report",
+                  {
+                    "id": "THE_ID",
+                    "message": "message1",
+                  },
+                ],
+                [
+                  "progress/report",
+                  {
+                    "id": "THE_ID",
+                    "increment": 50,
+                  },
+                ],
+                [
+                  "progress/report",
+                  {
+                    "id": "THE_ID",
+                    "increment": 50,
+                  },
+                ],
+                [
+                  "progress/end",
+                  {
+                    "id": "THE_ID",
+                  },
+                ],
+              ]
+            `)
+        })
+        it('progress can be cancelled', async () => {
+            const disposable = client.progressStartEvents.event(params => {
+                if (params.options.title === 'testing/progressCancelation') {
+                    client.notify('progress/cancel', { id: params.id })
+                }
+            })
+            try {
+                const { result } = await client.request('testing/progressCancelation', { title: 'Leona' })
+                expect(result).toStrictEqual("request with title 'Leona' cancelled")
+            } finally {
+                disposable.dispose()
+            }
+        })
+    })
+
     describe('RateLimitedAgent', () => {
         const rateLimitedClient = new TestClient('rateLimitedClient', process.env.SRC_ACCESS_TOKEN_WITH_RATE_LIMIT)
         // Initialize inside beforeAll so that subsequent tests are skipped if initialization fails.

agent/src/vscode-shim.ts

@@ -2,6 +2,7 @@
 import { execSync } from 'child_process'
 import path from 'path'
 
+import * as uuid from 'uuid'
 import type * as vscode from 'vscode'
 
 // <VERY IMPORTANT - PLEASE READ>
@@ -31,6 +32,7 @@ import {
     ExtensionKind,
     FileType,
     LogLevel,
+    ProgressLocation,
     Range,
     StatusBarAlignment,
     UIKind,
@@ -414,6 +416,8 @@ export function setCreateWebviewPanel(newCreateWebviewPanel: typeof vscode.windo
     shimmedCreateWebviewPanel = newCreateWebviewPanel
 }
 
+export const progressBars = new Map<string, CancellationTokenSource>()
+
 const _window: Partial<typeof vscode.window> = {
     createTreeView: () => defaultTreeView,
     tabGroups,
@@ -443,13 +447,43 @@ const _window: Partial<typeof vscode.window> = {
     registerWebviewViewProvider: () => emptyDisposable,
     createStatusBarItem: () => statusBarItem,
     visibleTextEditors,
-    withProgress: async (_, handler) => {
+    withProgress: async (options, handler) => {
+        const progressClient = clientInfo?.capabilities?.progressBars === 'enabled' ? agent : undefined
+        const id = uuid.v4()
+        const tokenSource = new CancellationTokenSource()
+        const token = tokenSource.token
+        progressBars.set(id, tokenSource)
+        token.onCancellationRequested(() => progressBars.delete(id))
+
+        if (progressClient) {
+            const location = typeof options.location === 'number' ? ProgressLocation[options.location] : undefined
+            const locationViewId = typeof options.location === 'object' ? options.location.viewId : undefined
+            progressClient.notify('progress/start', {
+                id,
+                options: { title: options.title, cancellable: options.cancellable, location, locationViewId },
+            })
+        }
         try {
-            const result = await handler({ report: () => {} }, new CancellationTokenSource().token)
+            const result = await handler(
+                {
+                    report: ({ message, increment }) => {
+                        if (progressClient && !token.isCancellationRequested) {
+                            progressClient.notify('progress/report', { id, message, increment })
+                        }
+                    },
+                },
+                token
+            )
             return result
         } catch (error) {
             console.error('window.withProgress: uncaught error', error)
             throw error
+        } finally {
+            tokenSource.dispose()
+            progressBars.delete(id)
+            if (progressClient) {
+                progressClient.notify('progress/end', { id })
+            }
         }
     },
     onDidChangeActiveTextEditor: onDidChangeActiveTextEditor.event,