Agent: make AgentWorkspaceDocuments more robust (#4279)

agent/src/AgentTextEditor.ts

@@ -0,0 +1,92 @@
+import { logDebug } from '@sourcegraph/cody-shared'
+import * as vscode from 'vscode'
+import type { AgentTextDocument } from './AgentTextDocument'
+import type { EditFunction } from './AgentWorkspaceDocuments'
+
+export class AgentTextEditor implements vscode.TextEditor {
+    constructor(
+        private readonly agentDocument: AgentTextDocument,
+        private readonly params?: { edit?: EditFunction }
+    ) {}
+    get document(): vscode.TextDocument {
+        return this.agentDocument
+    }
+    get selection(): vscode.Selection {
+        const protocolSelection = this.agentDocument.protocolDocument.selection
+        const selection: vscode.Selection = protocolSelection
+            ? new vscode.Selection(
+                  new vscode.Position(protocolSelection.start.line, protocolSelection.start.character),
+                  new vscode.Position(protocolSelection.end.line, protocolSelection.end.character)
+              )
+            : // Default to putting the cursor at the start of the file.
+              new vscode.Selection(new vscode.Position(0, 0), new vscode.Position(0, 0))
+        return selection
+    }
+    get selections(): readonly vscode.Selection[] {
+        return [this.selection]
+    }
+    get visibleRanges(): readonly vscode.Range[] {
+        const protocolVisibleRange = this.agentDocument.protocolDocument.visibleRange
+        const visibleRange = protocolVisibleRange
+            ? new vscode.Selection(
+                  new vscode.Position(
+                      protocolVisibleRange.start.line,
+                      protocolVisibleRange.start.character
+                  ),
+                  new vscode.Position(protocolVisibleRange.end.line, protocolVisibleRange.end.character)
+              )
+            : this.selection
+        return [visibleRange]
+    }
+    get options(): vscode.TextEditorOptions {
+        return {
+            cursorStyle: undefined,
+            insertSpaces: undefined,
+            lineNumbers: undefined,
+            // TODO: fix tabSize
+            tabSize: 2,
+        }
+    }
+    viewColumn = vscode.ViewColumn.Active
+
+    // IMPORTANT(olafurpg): `edit` must be defined as a fat arrow. The tests
+    // fail if it's defined as a normal class method.
+    edit = (
+        callback: (editBuilder: vscode.TextEditorEdit) => void,
+        options?: { readonly undoStopBefore: boolean; readonly undoStopAfter: boolean } | undefined
+    ): Promise<boolean> => {
+        if (this.params?.edit) {
+            return this.params.edit(this.agentDocument.uri, callback, options)
+        }
+        logDebug('AgentTextEditor:edit()', 'not supported')
+        return Promise.resolve(false)
+    }
+    insertSnippet(
+        snippet: vscode.SnippetString,
+        location?:
+            | vscode.Range
+            | vscode.Position
+            | readonly vscode.Range[]
+            | readonly vscode.Position[]
+            | undefined,
+        options?: { readonly undoStopBefore: boolean; readonly undoStopAfter: boolean } | undefined
+    ): Thenable<boolean> {
+        // Do nothing, for now.
+        return Promise.resolve(true)
+    }
+    setDecorations(
+        decorationType: vscode.TextEditorDecorationType,
+        rangesOrOptions: readonly vscode.Range[] | readonly vscode.DecorationOptions[]
+    ): void {
+        // Do nothing, for now
+    }
+    revealRange(range: vscode.Range, revealType?: vscode.TextEditorRevealType | undefined): void {
+        // Do nothing, for now.
+    }
+    show(column?: vscode.ViewColumn | undefined): void {
+        // Do nothing, for now.
+    }
+    hide(): void {
+        // Do nothing, for now.
+    }
+}

agent/src/AgentWorkspaceDocuments.test.ts

@@ -0,0 +1,119 @@
+import { beforeEach, describe, expect, it } from 'vitest'
+import * as vscode from 'vscode'
+import { ProtocolTextDocumentWithUri } from '../../vscode/src/jsonrpc/TextDocumentWithUri'
+import { AgentWorkspaceDocuments } from './AgentWorkspaceDocuments'
+
+describe('AgentWorkspaceDocuments', () => {
+    let documents: AgentWorkspaceDocuments
+    beforeEach(() => {
+        documents = new AgentWorkspaceDocuments({})
+    })
+    const uri = vscode.Uri.parse('file:///foo.txt')
+
+    it('singleton document', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'hello' })
+        )
+        expect(document.getText()).toBe('hello')
+        const document2 = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'goodbye' })
+        )
+        // Regardless of when you got the reference to the document, `getText()`
+        // always reflects the latest value.
+        expect(document.getText()).toBe('goodbye')
+        expect(document2.getText()).toBe('goodbye')
+    })
+
+    it('null content', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: 'hello' })
+        )
+        expect(document.getText()).toBe('hello')
+        expect(documents.getDocument(uri)?.getText()).toBe('hello')
+
+        const document2 = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                contentChanges: null as any,
+                content: null as any,
+                visibleRange: null as any,
+                selection: null as any,
+            })
+        )
+        expect(document2.getText()).toBe('hello')
+        expect(document2.protocolDocument.contentChanges).toBeUndefined()
+        expect(document2.protocolDocument.selection).toBeUndefined()
+        expect(document2.protocolDocument.visibleRange).toBeUndefined()
+    })
+
+    it('incremental sync', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, { content: ['abc', 'def', 'ghi'].join('\n') })
+        )
+        expect(document.getText()).toBe('abc\ndef\nghi')
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                contentChanges: [
+                    {
+                        range: { start: { line: 0, character: 0 }, end: { line: 0, character: 1 } },
+                        text: 'x',
+                    },
+                    {
+                        range: { start: { line: 1, character: 1 }, end: { line: 1, character: 2 } },
+                        text: 'y',
+                    },
+                    {
+                        range: { start: { line: 2, character: 2 }, end: { line: 2, character: 3 } },
+                        text: 'z',
+                    },
+                ],
+            })
+        )
+        expect(document.getText()).toBe('xbc\ndyf\nghz')
+    })
+
+    it('selection', () => {
+        const document = documents.loadAndUpdateDocument(ProtocolTextDocumentWithUri.from(uri, {}))
+        const editor = documents.newTextEditor(document)
+        expect(editor.selection).toStrictEqual(new vscode.Selection(0, 0, 0, 0))
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                content: 'hello\ngoodbye\nworld\nsayonara\n',
+                selection: { start: { line: 0, character: 0 }, end: { line: 1, character: 5 } },
+            })
+        )
+        const expectedSelection = new vscode.Selection(0, 0, 1, 5)
+        expect(editor.selection).toStrictEqual(expectedSelection)
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                content: 'something\nis\nhappening',
+                visibleRange: undefined,
+            })
+        )
+        expect(editor.selection).toStrictEqual(expectedSelection)
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                selection: { start: { line: 1, character: 1 }, end: { line: 2, character: 3 } },
+            })
+        )
+        expect(editor.selection).toStrictEqual(new vscode.Selection(1, 1, 2, 3))
+    })
+
+    it('visibleRanges', () => {
+        const document = documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                content: 'hello\ngoodbye\nworld\nsayonara\n',
+                visibleRange: { start: { line: 0, character: 0 }, end: { line: 1, character: 5 } },
+            })
+        )
+        const editor = documents.newTextEditor(document)
+        const expectedSelection = new vscode.Selection(0, 0, 1, 5)
+        expect(editor.visibleRanges).toStrictEqual([expectedSelection])
+        documents.loadAndUpdateDocument(
+            ProtocolTextDocumentWithUri.from(uri, {
+                content: 'something\nis\nhappening',
+                visibleRange: undefined,
+            })
+        )
+        expect(editor.visibleRanges).toStrictEqual([expectedSelection])
+    })
+})