From 09686103a65a0e6aa819d6862fbe6fc9197b6095 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=93lafur=20P=C3=A1ll=20Geirsson?= Date: Fri, 24 May 2024 09:16:49 +0200 Subject: [PATCH] Agent: make AgentWorkspaceDocuments more robust (#4279) --- agent/src/AgentTextEditor.ts | 92 ++++++++++++++++ agent/src/AgentWorkspaceDocuments.test.ts | 119 +++++++++++++++++++++ agent/src/AgentWorkspaceDocuments.ts | 123 ++++++++++------------ agent/src/TestClient.ts | 2 +- 4 files changed, 265 insertions(+), 71 deletions(-) create mode 100644 agent/src/AgentTextEditor.ts create mode 100644 agent/src/AgentWorkspaceDocuments.test.ts diff --git a/agent/src/AgentTextEditor.ts b/agent/src/AgentTextEditor.ts new file mode 100644 index 000000000000..1ec6a9dbdb06 --- /dev/null +++ b/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 => { + 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 { + // 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. + } +} diff --git a/agent/src/AgentWorkspaceDocuments.test.ts b/agent/src/AgentWorkspaceDocuments.test.ts new file mode 100644 index 000000000000..e43647b22638 --- /dev/null +++ b/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]) + }) +}) diff --git a/agent/src/AgentWorkspaceDocuments.ts b/agent/src/AgentWorkspaceDocuments.ts index 1f7aa7a798fa..f7d662651282 100644 --- a/agent/src/AgentWorkspaceDocuments.ts +++ b/agent/src/AgentWorkspaceDocuments.ts @@ -4,16 +4,17 @@ import * as vscode from 'vscode' import { ProtocolTextDocumentWithUri } from '../../vscode/src/jsonrpc/TextDocumentWithUri' -import { logDebug, logError } from '@sourcegraph/cody-shared' +import { logError } from '@sourcegraph/cody-shared' import { doesFileExist } from '../../vscode/src/commands/utils/workspace-files' import { resetActiveEditor } from '../../vscode/src/editor/active-editor' import { AgentTextDocument } from './AgentTextDocument' +import { AgentTextEditor } from './AgentTextEditor' import { applyContentChanges } from './applyContentChanges' import { calculateContentChanges } from './calculateContentChanges' import { clearArray } from './clearArray' import * as vscode_shim from './vscode-shim' -type EditFunction = ( +export type EditFunction = ( uri: vscode.Uri, callback: (editBuilder: vscode.TextEditorEdit) => void, options?: { readonly undoStopBefore: boolean; readonly undoStopAfter: boolean } @@ -27,7 +28,10 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { constructor(private params?: { edit?: EditFunction }) {} // Keys are `vscode.Uri.toString()` formatted. We don't use `vscode.Uri` as // keys because hashcode/equals behave unreliably. - private readonly agentDocuments: Map = new Map() + private readonly agentDocuments: Map< + string, + { document: AgentTextDocument; editor: AgentTextEditor } + > = new Map() public workspaceRootUri: vscode.Uri | undefined public activeDocumentFilePath: vscode.Uri | null = null @@ -36,44 +40,69 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { return this.loadAndUpdateDocument(ProtocolTextDocumentWithUri.from(uri)) } public loadAndUpdateDocument(document: ProtocolTextDocumentWithUri): AgentTextDocument { - return this.loadAndUpdateDocumentWithChanges(document).document + return this.loadAndUpdate(document).document } - public loadAndUpdateDocumentWithChanges(document: ProtocolTextDocumentWithUri): { + public loadAndUpdate(document: ProtocolTextDocumentWithUri): { document: AgentTextDocument + editor: AgentTextEditor contentChanges: vscode.TextDocumentContentChangeEvent[] } { - const fromCache = this.agentDocuments.get(document.underlying.uri) - if (!fromCache) { - return { document: new AgentTextDocument(document), contentChanges: [] } + const cached = this.agentDocuments.get(document.underlying.uri) + if (!cached) { + const result = new AgentTextDocument(document) + const editor = new AgentTextEditor(result, this.params) + this.agentDocuments.set(document.underlying.uri, { document: result, editor }) + return { document: result, editor, contentChanges: [] } } + const { document: fromCache } = cached + // We have seen this document before, which means we mutate the existing + // document to reflect the latest chagnes. For each URI, we keep a + // singleton document so that `AgentTextDocument.getText()` always + // reflects the latest value. const contentChanges: vscode.TextDocumentContentChangeEvent[] = [] - if (document.contentChanges && document.contentChanges.length > 0) { + // Incremental document sync. const changes = applyContentChanges(fromCache, document.contentChanges) contentChanges.push(...changes.contentChanges) document.underlying.content = changes.newText - } else if (document.content !== undefined && document.content != null) { + } else if (typeof document.content === 'string') { + // Full document sync. for (const change of calculateContentChanges(fromCache, document.content)) { contentChanges.push(change) } - } - - if (document.content === undefined) { + } else { + // No document sync. Use content from last update. document.underlying.content = fromCache.getText() } - if (document.selection === undefined) { + if (!document.selection) { + // No changes to the selection, populate from cache document.underlying.selection = fromCache.protocolDocument.selection } - if (document.visibleRange === undefined) { + if (!document.visibleRange) { + // No changes to the visible range, populate from cache document.underlying.visibleRange = fromCache.protocolDocument.visibleRange } + // The client may send null values that we convert to undefined here. + if (document.content === null) { + document.underlying.content = undefined + } + if (document.contentChanges === null) { + document.underlying.contentChanges = undefined + } + if (document.selection === null) { + document.underlying.selection = undefined + } + if (document.visibleRange === null) { + document.underlying.visibleRange = undefined + } + fromCache.update(document) - return { document: fromCache, contentChanges } + return { document: fromCache, editor: cached.editor, contentChanges } } public setActiveTextEditor(textEditor: vscode.TextEditor): void { @@ -87,15 +116,15 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { } public allDocuments(): AgentTextDocument[] { - return [...this.agentDocuments.values()] + return [...this.agentDocuments.values()].map(value => value.document) } public getDocument(uri: vscode.Uri): AgentTextDocument | undefined { - return this.agentDocuments.get(uri.toString()) + return this.agentDocuments.get(uri.toString())?.document } public getDocumentFromUriString(uriString: string): AgentTextDocument | undefined { - return this.agentDocuments.get(uriString) + return this.agentDocuments.get(uriString)?.document } public loadDocument(document: ProtocolTextDocumentWithUri): AgentTextDocument { @@ -105,9 +134,7 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { document: AgentTextDocument contentChanges: vscode.TextDocumentContentChangeEvent[] } { - const { document: agentDocument, contentChanges } = - this.loadAndUpdateDocumentWithChanges(document) - this.agentDocuments.set(document.underlying.uri, agentDocument) + const { document: agentDocument, contentChanges } = this.loadAndUpdate(document) const tabs: vscode.Tab[] = [] for (const uri of this.allUris()) { @@ -195,53 +222,9 @@ export class AgentWorkspaceDocuments implements vscode_shim.WorkspaceDocuments { } public newTextEditor(document: AgentTextDocument): vscode.TextEditor { - const protocolSelection = document.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) - ) - : new vscode.Selection(new vscode.Position(0, 0), new vscode.Position(0, 0)) - - const protocolVisibleRange = document.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) - ) - : selection - - 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: (callback, options) => { - if (this.params?.edit) { - return this.params.edit(document.uri, callback, options) - } - logDebug('AgentTextEditor:edit()', 'not supported') - return Promise.resolve(false) - }, - insertSnippet: () => Promise.resolve(true), - revealRange: () => {}, // TODO: implement this for inline edit commands? - options: { - cursorStyle: undefined, - insertSpaces: undefined, - lineNumbers: undefined, - // TODO: fix tabSize - tabSize: 2, - }, - setDecorations: () => {}, - viewColumn: vscode.ViewColumn.Active, - visibleRanges: [visibleRange], - show: () => {}, - hide: () => {}, - } + return ( + this.agentDocuments.get(document.protocolDocument.underlying.uri)?.editor ?? + new AgentTextEditor(document, this.params) + ) } } diff --git a/agent/src/TestClient.ts b/agent/src/TestClient.ts index 5c144c769e78..e479c01739f5 100644 --- a/agent/src/TestClient.ts +++ b/agent/src/TestClient.ts @@ -177,7 +177,7 @@ export class TestClient extends MessageHandler { public progressIDs = new Map() public progressStartEvents = new vscode.EventEmitter() public readonly name: string - public workspace = new AgentWorkspaceDocuments() + public workspace = new AgentWorkspaceDocuments({}) public workspaceEditParams: WorkspaceEditParams[] = [] public textDocumentEditParams: TextDocumentEditParams[] = []