Agent: add support for autocomplete (#624)

This PR adds a new `autocomplete/execute` endpoint to the agent
protocol. To implement this endpoint, the agent now directly depend on
the VS Code extension by providing a shim for all imports against the
`'vscode'` module. By using a shim, the agent can launch the VS Code
extension in a headless environment and capture the instance o the
completion provider when it's registered.

This PR doesn't implement all of VS Code APIs yet. It only implements
that necessary parts to get autocomplete working.


## Test plan

See new `index.test.ts`
<!-- Required. See
https://docs.sourcegraph.com/dev/background-information/testing_principles.
-->

---------

Signed-off-by: Stephen Gutekanst <[email protected]>
Co-authored-by: Chris Warwick <[email protected]>
Co-authored-by: Stephen Gutekanst <[email protected]>

agent/package.json

@@ -12,10 +12,10 @@
   "main": "src/index.ts",
   "sideEffects": false,
   "scripts": {
-    "build": "esbuild ./src/index.ts --bundle --outfile=dist/agent.js --format=cjs --platform=node",
-    "build-minify": "pnpm build --minify",
+    "build": "node ./src/esbuild.mjs",
+    "agent": "pnpm run build && node dist/index.js",
     "build-ts": "tsc --build",
-    "build-agent-binaries": "pnpm run build && pkg -t node16-linux-arm64,node16-linux-x64,node16-macos-arm64,node16-macos-x64,node16-win-x64 dist/agent.js --out-path ${AGENT_EXECUTABLE_TARGET_DIRECTORY:-dist}",
+    "build-agent-binaries": "pnpm run build && cp dist/index.js dist/agent.js && pkg -t node16-linux-arm64,node16-linux-x64,node16-macos-arm64,node16-macos-x64,node16-win-x64 dist/agent.js --out-path ${AGENT_EXECUTABLE_TARGET_DIRECTORY:-dist}",
     "lint": "pnpm run lint:js",
     "lint:js": "eslint --cache '**/*.[tj]s?(x)'",
     "test": "pnpm run build && vitest"
@@ -25,6 +25,9 @@
     "vscode-uri": "^3.0.7"
   },
   "devDependencies": {
-    "pkg": "^5.8.1"
+    "pkg": "^5.8.1",
+    "@types/vscode": "^1.80.0",
+    "esbuild-plugin-alias-path": "^2.0.2",
+    "esbuild": "^0.18.19"
   }
 }

agent/src/AgentTabGroups.ts

@@ -0,0 +1,13 @@
+import type * as vscode from 'vscode'
+
+import { emptyEvent } from '../../vscode/src/testutils/mocks'
+
+export class AgentTabGroups implements vscode.TabGroups {
+    public all: vscode.TabGroup[] = []
+    public activeTabGroup: vscode.TabGroup = { activeTab: undefined, isActive: true, tabs: [], viewColumn: 1 }
+    public onDidChangeTabGroups: vscode.Event<vscode.TabGroupChangeEvent> = emptyEvent()
+    public onDidChangeTabs: vscode.Event<vscode.TabChangeEvent> = emptyEvent()
+    public close(): Thenable<boolean> {
+        throw new Error('Method not implemented.')
+    }
+}

agent/src/AgentTextDocument.test.ts

@@ -0,0 +1,74 @@
+import assert from 'assert'
+
+import { describe, it } from 'vitest'
+import * as vscode from 'vscode'
+
+import { AgentTextDocument } from './AgentTextDocument'
+
+describe('AgentTextDocument', () => {
+    const basic = new AgentTextDocument({ filePath: 'foo', content: 'a\nb\n' })
+    const basicCrlf = new AgentTextDocument({ filePath: 'foo', content: 'a\r\nb\r\n' })
+    const emptyLine = new AgentTextDocument({ filePath: 'foo', content: 'a\n\n' })
+    const noEndOfFileNewline = new AgentTextDocument({ filePath: 'foo', content: 'a\nb' })
+    const emptyFirstLine = new AgentTextDocument({ filePath: 'foo', content: '\nb' })
+    const emptyFirstLineCrlf = new AgentTextDocument({ filePath: 'foo', content: '\r\nb' })
+    const noIndentation = new AgentTextDocument({ filePath: 'foo', content: 'sss\n' })
+    const indentation = new AgentTextDocument({ filePath: 'foo', content: '  a\n' })
+    const indentationTab = new AgentTextDocument({ filePath: 'foo', content: '\t\tab\n' })
+
+    it('getText(Range)', () => {
+        assert.equal(basic.getText(new vscode.Range(0, 0, 0, 1)), 'a')
+        assert.equal(basic.getText(new vscode.Range(0, 0, 1, 1)), 'a\nb')
+        assert.equal(basic.getText(new vscode.Range(2, 0, 2, 10)), '')
+        assert.equal(basic.getText(new vscode.Range(0, 0, 2, 3)), 'a\nb\n')
+    })
+
+    it('lineCount()', () => {
+        assert.equal(basic.lineCount, 2)
+        assert.equal(basicCrlf.lineCount, 2)
+        assert.equal(emptyFirstLine.lineCount, 2)
+        assert.equal(noEndOfFileNewline.lineCount, 2)
+        assert.equal(emptyFirstLine.lineCount, 2)
+        assert.equal(emptyFirstLineCrlf.lineCount, 2)
+    })
+
+    it('positionAt()', () => {
+        assert.deepEqual(basic.positionAt(0), new vscode.Position(0, 0))
+    })
+
+    it('lineAt()', () => {
+        assert.equal(basic.getText(basic.lineAt(1).range), 'b')
+        assert.equal(basic.getText(basic.lineAt(1).rangeIncludingLineBreak), 'b\n')
+        assert.equal(basic.getText(basic.lineAt(2).range), '')
+        assert.equal(basic.getText(basic.lineAt(2).rangeIncludingLineBreak), '')
+
+        assert.equal(basicCrlf.getText(basic.lineAt(1).range), 'b')
+        assert.equal(basicCrlf.getText(basicCrlf.lineAt(1).rangeIncludingLineBreak), 'b\r\n')
+        assert.equal(basicCrlf.getText(basic.lineAt(2).range), '')
+        assert.equal(basicCrlf.getText(basic.lineAt(2).rangeIncludingLineBreak), '')
+
+        assert.equal(emptyLine.getText(emptyLine.lineAt(0).range), 'a')
+        assert.equal(emptyLine.getText(emptyLine.lineAt(0).rangeIncludingLineBreak), 'a\n')
+        assert.equal(emptyLine.getText(emptyLine.lineAt(1).range), '')
+        assert.equal(emptyLine.getText(emptyLine.lineAt(1).rangeIncludingLineBreak), '\n')
+
+        assert.equal(noEndOfFileNewline.getText(noEndOfFileNewline.lineAt(1).range), 'b')
+        assert.equal(noEndOfFileNewline.getText(noEndOfFileNewline.lineAt(1).rangeIncludingLineBreak), 'b')
+
+        assert.equal(emptyFirstLine.getText(emptyFirstLine.lineAt(0).range), '')
+        assert.equal(emptyFirstLine.getText(emptyFirstLine.lineAt(0).rangeIncludingLineBreak), '\n')
+        assert.equal(emptyFirstLine.getText(emptyFirstLine.lineAt(1).range), 'b')
+        assert.equal(emptyFirstLine.getText(emptyFirstLine.lineAt(1).rangeIncludingLineBreak), 'b')
+
+        assert.equal(emptyFirstLineCrlf.getText(emptyFirstLineCrlf.lineAt(0).range), '')
+        assert.equal(emptyFirstLineCrlf.getText(emptyFirstLineCrlf.lineAt(0).rangeIncludingLineBreak), '\r\n')
+        assert.equal(emptyFirstLineCrlf.getText(emptyFirstLineCrlf.lineAt(1).range), 'b')
+        assert.equal(emptyFirstLineCrlf.getText(emptyFirstLineCrlf.lineAt(1).rangeIncludingLineBreak), 'b')
+    })
+
+    it('lineAt().firstNonWhitespaceCharacterIndex', () => {
+        assert.equal(noIndentation.lineAt(0).firstNonWhitespaceCharacterIndex, 0)
+        assert.equal(indentation.lineAt(0).firstNonWhitespaceCharacterIndex, 2)
+        assert.equal(indentationTab.lineAt(0).firstNonWhitespaceCharacterIndex, 2)
+    })
+})

agent/src/AgentTextDocument.ts

@@ -0,0 +1,88 @@
+import * as vscode from 'vscode'
+
+import { DocumentOffsets } from './offsets'
+import { TextDocument } from './protocol'
+import * as vscode_shim from './vscode-shim'
+
+// TODO: implement with vscode-languageserver-textdocument The reason we don't
+// use vscode-languageserver-textdocument is because it doesn't implement all
+// the properties/functions that vscode.TextDocument has. For example, lineAt is
+// missing in vscode-languageserver-textdocument
+export class AgentTextDocument implements vscode.TextDocument {
+    constructor(public readonly textDocument: TextDocument) {
+        this.content = textDocument.content ?? ''
+        this.uri = vscode_shim.Uri.from({ scheme: 'file', path: textDocument.filePath })
+        this.fileName = textDocument.filePath
+        this.isUntitled = false
+        this.languageId = this.fileName.split('.').splice(-1)[0]
+        this.offsets = new DocumentOffsets(textDocument)
+        this.lineCount = this.offsets.lineCount()
+    }
+    private readonly content: string
+    private readonly offsets: DocumentOffsets
+    public readonly uri: vscode.Uri
+    public readonly fileName: string
+    public readonly lineCount: number
+    public readonly isUntitled: boolean
+    public readonly languageId: string
+
+    public readonly version: number = 0
+    public readonly isDirty: boolean = false
+    public readonly isClosed: boolean = false
+    public save(): Thenable<boolean> {
+        throw new Error('Method not implemented.')
+    }
+    public readonly eol: vscode.EndOfLine = vscode_shim.EndOfLine.LF
+    public lineAt(position: vscode.Position | number): vscode.TextLine {
+        const line = typeof position === 'number' ? position : position.line
+        const text = this.getText(
+            new vscode_shim.Range(
+                new vscode_shim.Position(line, 0),
+                new vscode_shim.Position(line, this.offsets.lineLengthExcludingNewline(line))
+            )
+        )
+        let firstNonWhitespaceCharacterIndex = 0
+        while (firstNonWhitespaceCharacterIndex < text.length && /\s/.test(text[firstNonWhitespaceCharacterIndex])) {
+            firstNonWhitespaceCharacterIndex++
+        }
+        return {
+            lineNumber: line,
+            firstNonWhitespaceCharacterIndex,
+            isEmptyOrWhitespace: firstNonWhitespaceCharacterIndex === text.length,
+            range: new vscode_shim.Range(
+                new vscode_shim.Position(line, 0),
+                new vscode_shim.Position(line, text.length)
+            ),
+            rangeIncludingLineBreak: new vscode_shim.Range(
+                new vscode_shim.Position(line, 0),
+                new vscode_shim.Position(line, text.length + this.offsets.newlineLength(line))
+            ),
+            text,
+        }
+    }
+    public offsetAt(position: vscode.Position): number {
+        return this.offsets.offset(position)
+    }
+    public positionAt(offset: number): vscode.Position {
+        const { line, character } = this.offsets.position(offset)
+        return new vscode_shim.Position(line, character)
+    }
+    public getText(range?: vscode.Range | undefined): string {
+        if (range === undefined) {
+            return this.content
+        }
+        const start = this.offsets.offset(range.start)
+        const end = this.offsets.offset(range.end)
+        const text = this.content.slice(start, end)
+        return text
+    }
+    public getWordRangeAtPosition(position: vscode.Position, regex?: RegExp | undefined): vscode.Range | undefined {
+        throw new Error('Method not implemented.')
+    }
+    public validateRange(range: vscode.Range): vscode.Range {
+        throw new Error('Method not implemented.')
+    }
+    public validatePosition(position: vscode.Position): vscode.Position {
+        throw new Error('Method not implemented.')
+    }
+}

agent/src/AgentTextEditor.ts

@@ -0,0 +1,42 @@
+import * as vscode from 'vscode'
+
+import { AgentTextDocument } from './AgentTextDocument'
+
+export function newTextEditor(document: AgentTextDocument): vscode.TextEditor {
+    const selection: vscode.Selection = document.textDocument.selection
+        ? new vscode.Selection(
+              new vscode.Position(
+                  document.textDocument.selection.start.line,
+                  document.textDocument.selection.start.character
+              ),
+              new vscode.Position(
+                  document.textDocument.selection.end.line,
+                  document.textDocument.selection.end.character
+              )
+          )
+        : new vscode.Selection(new vscode.Position(0, 0), new vscode.Position(0, 0))
+
+    return {
+        // Looking at the implementation of the extension, we only need
+        // to provide `document` but we do a best effort to shim the
+        // rest of the `TextEditor` properties.
+        document,
+        selection,
+        selections: [selection],
+        edit: () => Promise.resolve(true),
+        insertSnippet: () => Promise.resolve(true),
+        revealRange: () => {},
+        options: {
+            cursorStyle: undefined,
+            insertSpaces: undefined,
+            lineNumbers: undefined,
+            // TODO: fix tabSize
+            tabSize: 2,
+        },
+        setDecorations: () => {},
+        viewColumn: vscode.ViewColumn.Active,
+        visibleRanges: [selection],
+        show: () => {},
+        hide: () => {},
+    }
+}

agent/src/AgentWorkspaceDocuments.ts

@@ -0,0 +1,73 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import type * as vscode from 'vscode'
+import { URI } from 'vscode-uri'
+
+import { AgentTextDocument } from './AgentTextDocument'
+import { newTextEditor } from './AgentTextEditor'
+import { TextDocument } from './protocol'
+import * as vscode_shim from './vscode-shim'
+
+export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments {
+    private readonly documents: Map<string, TextDocument> = new Map()
+    public workspaceRootUri: URI | null = null
+    public activeDocumentFilePath: string | null = null
+    public loadedDocument(document: TextDocument): TextDocument {
+        const fromCache = this.documents.get(document.filePath)
+        if (document.content === undefined) {
+            document.content = fromCache?.content
+        }
+        if (document.selection === undefined) {
+            document.selection = fromCache?.selection
+        }
+        return document
+    }
+    public agentTextDocument(document: TextDocument): AgentTextDocument {
+        return new AgentTextDocument(this.loadedDocument(document))
+    }
+
+    public allFilePaths(): string[] {
+        return [...this.documents.keys()]
+    }
+    public allDocuments(): TextDocument[] {
+        return [...this.documents.values()]
+    }
+    public getDocument(filePath: string): TextDocument | undefined {
+        return this.documents.get(filePath)
+    }
+    public setDocument(document: TextDocument): void {
+        this.documents.set(document.filePath, this.loadedDocument(document))
+        const tabs: readonly vscode.Tab[] = this.allFilePaths().map(filePath => this.vscodeTab(filePath))
+        vscode_shim.tabGroups.all = [
+            {
+                tabs,
+                isActive: true,
+                activeTab: this.vscodeTab(this.activeDocumentFilePath ?? ''),
+                viewColumn: vscode_shim.ViewColumn.Active,
+            },
+        ]
+
+        while (vscode_shim.visibleTextEditors.length > 0) {
+            vscode_shim.visibleTextEditors.pop()
+        }
+        for (const document of this.allDocuments()) {
+            vscode_shim.visibleTextEditors.push(newTextEditor(this.agentTextDocument(document)))
+        }
+    }
+    public deleteDocument(filePath: string): void {
+        this.documents.delete(filePath)
+    }
+    private vscodeTab(filePath: string): vscode.Tab {
+        return {
+            input: {
+                uri: filePath,
+            },
+            label: 'label',
+            group: { activeTab: undefined, isActive: false, tabs: [], viewColumn: -1 },
+        } as any
+    }
+
+    public openTextDocument(filePath: string): Promise<vscode.TextDocument> {
+        return Promise.resolve(this.agentTextDocument({ filePath }))
+    }
+}