From 588501fd635c7a84d7eca1ccb5690d8a9d82980b Mon Sep 17 00:00:00 2001 From: Hitesh Sagtani Date: Tue, 24 Sep 2024 07:04:53 +0530 Subject: [PATCH] Context: add three new experimental retriever strategies (#5494) # Context The PR adds three major context sources useful for autocomplete and next-edit-suggestion. ### Recent Copy retriever. Developers often copy/cut and paste context and before pasting we show autocomplete suggestions to the user. The PR leverages the copied content on the clipboard by the user as a context source. I wasn't able to find any vscode api event exposed which triggers when user `copy` or `cut` text in the editor. This seems like an open issue: https://github.com/microsoft/vscode/issues/30066 Another alternative I think of is to use a keybinding of `Ctrl+x` and `Ctrl+c`, but not fully sure about its implications, since this is one of the most common shortcuts. As a workaround, The way current PR accomplishes the same is: 1. Tracks the selection made by the user in the last `1 minutes` and keeps tracks of upto `100` most recent selections. 3. At the time of retrieval, checks if the current clipboard content matches the selection items in the list. ### Recent View Ports This context source captures and utilizes the recently viewed portions of code in the editor. It keeps track of the visible areas of code that the developer has scrolled through or focused on within a specified time frame. ### Diagnostics The Diagnostics context source leverages the diagnostic information provided by VS Code and language servers. It collects and utilizes information about errors, for a file as a context source. ## Test plan 1. Automated test - Added CI tests for each of the retrievers 2. Manual test - Override the setting `"cody.autocomplete.experimental.graphContext": "recent-copy"` in vscode settings. - Observe the context events using `Autocomplete Trace View` --- .../completions/completion-provider-config.ts | 3 + .../completions/context/context-strategy.ts | 25 +- .../diagnostics-retriever.test.ts | 419 ++++++++++++++++++ .../diagnostics-retriever.ts | 174 ++++++++ .../recent-user-actions/recent-copy.test.ts | 160 +++++++ .../recent-user-actions/recent-copy.ts | 127 ++++++ .../recent-edits-retriever.test.ts | 0 .../recent-edits-retriever.ts | 2 +- .../recent-view-port.test.ts | 179 ++++++++ .../recent-user-actions/recent-view-port.ts | 113 +++++ vscode/src/completions/context/utils.ts | 3 + .../supercompletions/get-supercompletion.ts | 2 +- .../supercompletion-provider.ts | 2 +- 13 files changed, 1204 insertions(+), 5 deletions(-) create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts rename vscode/src/completions/context/retrievers/{recent-edits => recent-user-actions}/recent-edits-retriever.test.ts (100%) rename vscode/src/completions/context/retrievers/{recent-edits => recent-user-actions}/recent-edits-retriever.ts (99%) create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts create mode 100644 vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts diff --git a/vscode/src/completions/completion-provider-config.ts b/vscode/src/completions/completion-provider-config.ts index 6e76491c153..e88a3b9e972 100644 --- a/vscode/src/completions/completion-provider-config.ts +++ b/vscode/src/completions/completion-provider-config.ts @@ -44,6 +44,9 @@ class CompletionProviderConfig { 'recent-edits-1m', 'recent-edits-5m', 'recent-edits-mixed', + 'recent-copy', + 'diagnostics', + 'recent-view-port', ] return resolvedConfig.pipe( mergeMap(({ configuration }) => { diff --git a/vscode/src/completions/context/context-strategy.ts b/vscode/src/completions/context/context-strategy.ts index a15ac0fdb5c..d9d2c5b3565 100644 --- a/vscode/src/completions/context/context-strategy.ts +++ b/vscode/src/completions/context/context-strategy.ts @@ -10,7 +10,10 @@ import type { ContextRetriever } from '../types' import type { BfgRetriever } from './retrievers/bfg/bfg-retriever' import { JaccardSimilarityRetriever } from './retrievers/jaccard-similarity/jaccard-similarity-retriever' import { LspLightRetriever } from './retrievers/lsp-light/lsp-light-retriever' -import { RecentEditsRetriever } from './retrievers/recent-edits/recent-edits-retriever' +import { DiagnosticsRetriever } from './retrievers/recent-user-actions/diagnostics-retriever' +import { RecentCopyRetriever } from './retrievers/recent-user-actions/recent-copy' +import { RecentEditsRetriever } from './retrievers/recent-user-actions/recent-edits-retriever' +import { RecentViewPortRetriever } from './retrievers/recent-user-actions/recent-view-port' import { loadTscRetriever } from './retrievers/tsc/load-tsc-retriever' export type ContextStrategy = @@ -26,6 +29,9 @@ export type ContextStrategy = | 'recent-edits-1m' | 'recent-edits-5m' | 'recent-edits-mixed' + | 'recent-copy' + | 'diagnostics' + | 'recent-view-port' export interface ContextStrategyFactory extends vscode.Disposable { getStrategy( @@ -82,6 +88,18 @@ export class DefaultContextStrategyFactory implements ContextStrategyFactory { this.localRetriever = new JaccardSimilarityRetriever() this.graphRetriever = new LspLightRetriever() break + case 'recent-copy': + this.localRetriever = new RecentCopyRetriever({ + maxAgeMs: 60 * 1000, + maxSelections: 100, + }) + break + case 'diagnostics': + this.localRetriever = new DiagnosticsRetriever() + break + case 'recent-view-port': + this.localRetriever = new RecentViewPortRetriever() + break case 'jaccard-similarity': this.localRetriever = new JaccardSimilarityRetriever() break @@ -148,7 +166,10 @@ export class DefaultContextStrategyFactory implements ContextStrategyFactory { case 'jaccard-similarity': case 'recent-edits': case 'recent-edits-1m': - case 'recent-edits-5m': { + case 'recent-edits-5m': + case 'recent-copy': + case 'diagnostics': + case 'recent-view-port': { if (this.localRetriever) { retrievers.push(this.localRetriever) } diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts new file mode 100644 index 00000000000..b3f160a4da3 --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts @@ -0,0 +1,419 @@ +import dedent from 'dedent' +import { XMLParser } from 'fast-xml-parser' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import * as vscode from 'vscode' +import { document } from '../../../test-helpers' +import { DiagnosticsRetriever } from './diagnostics-retriever' + +describe('DiagnosticsRetriever', () => { + let retriever: DiagnosticsRetriever + let parser: XMLParser + + beforeEach(() => { + vi.useFakeTimers() + retriever = new DiagnosticsRetriever() + parser = new XMLParser() + }) + + afterEach(() => { + retriever.dispose() + }) + + // Helper function to reduce repetition in tests + const testDiagnostics = async ( + testDocument: vscode.TextDocument, + diagnostics: vscode.Diagnostic[], + position: vscode.Position, + expectedSnippetCount: number, + expectedMessageSnapshot: string + ) => { + const snippets = await retriever.getDiagnosticsPromptFromInformation( + testDocument, + position, + diagnostics + ) + expect(snippets).toHaveLength(expectedSnippetCount) + const message = parser.parse(snippets[0].content) + expect(message).toBeDefined() + expect(message.diagnostic).toBeDefined() + expect(message.diagnostic.message).toMatchInlineSnapshot(expectedMessageSnapshot) + return { snippets, message } + } + + // Helper function to create a diagnostic + const createDiagnostic = ( + severity: vscode.DiagnosticSeverity, + range: vscode.Range, + message: string, + source = 'ts', + relatedInformation?: vscode.DiagnosticRelatedInformation[] + ): vscode.Diagnostic => ({ + severity, + range, + message, + source, + relatedInformation, + }) + + it('should retrieve diagnostics for a given position', async () => { + const testDocument = document( + dedent` + function foo() { + console.log('foo') + } + `, + 'typescript' + ) + const diagnostic = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 16, 1, 21), + "Type 'string' is not assignable to type 'number'." + ), + ] + const position = new vscode.Position(1, 16) + + await testDiagnostics( + testDocument, + diagnostic, + position, + 1, + ` + "function foo() { + console.log('foo') + ^^^^^ Type 'string' is not assignable to type 'number'. + }" + ` + ) + }) + + it('should retrieve diagnostics on multiple lines', async () => { + const testDocument = document( + dedent` + function multiLineErrors() { + const x: number = "string"; + const y: string = 42; + const z = x + y; + } + `, + 'typescript' + ) + const diagnostics = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 24, 1, 32), + "Type 'string' is not assignable to type 'number'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(2, 24, 2, 26), + "Type 'number' is not assignable to type 'string'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(3, 18, 3, 23), + "The '+' operator cannot be applied to types 'number' and 'string'." + ), + ] + const position = new vscode.Position(1, 0) + + const { snippets } = await testDiagnostics( + testDocument, + diagnostics, + position, + 3, + ` + "function multiLineErrors() { + const x: number = "string"; + ^^^^^^^^ Type 'string' is not assignable to type 'number'. + const y: string = 42; + const z = x + y; + }" + ` + ) + + expect(parser.parse(snippets[1].content).diagnostic.message).toMatchInlineSnapshot(` + "function multiLineErrors() { + const x: number = "string"; + const y: string = 42; + ^^ Type 'number' is not assignable to type 'string'. + const z = x + y; + }" + `) + expect(parser.parse(snippets[2].content).diagnostic.message).toMatchInlineSnapshot(` + "function multiLineErrors() { + const x: number = "string"; + const y: string = 42; + const z = x + y; + ^^^ The '+' operator cannot be applied to types 'number' and 'string'. + }" + `) + }) + + it('should handle multiple diagnostics on the same line', async () => { + const testDocument = document( + dedent` + function bar(x: number, y: string) { + return x + y; + } + `, + 'typescript' + ) + const diagnostics = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 11, 1, 12), + "The '+' operator cannot be applied to types 'number' and 'string'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 14, 1, 15), + "Implicit conversion of 'string' to 'number' may cause unexpected behavior." + ), + ] + const position = new vscode.Position(1, 11) + + await testDiagnostics( + testDocument, + diagnostics, + position, + 1, + ` + "function bar(x: number, y: string) { + return x + y; + ^ The '+' operator cannot be applied to types 'number' and 'string'. + ^ Implicit conversion of 'string' to 'number' may cause unexpected behavior. + }" + ` + ) + }) + + it('should filter out warning diagnostics', async () => { + const testDocument = document( + dedent` + function bar(x: number, y: string) { + return x + y; + } + `, + 'typescript' + ) + const diagnostics = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 11, 1, 12), + "The '+' operator cannot be applied to types 'number' and 'string'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Warning, + new vscode.Range(1, 14, 1, 15), + "Implicit conversion of 'string' to 'number' may cause unexpected behavior." + ), + ] + const position = new vscode.Position(1, 11) + + await testDiagnostics( + testDocument, + diagnostics, + position, + 1, + ` + "function bar(x: number, y: string) { + return x + y; + ^ The '+' operator cannot be applied to types 'number' and 'string'. + }" + ` + ) + }) + + it('should handle diagnostics at the end of the file', async () => { + const testDocument = document( + dedent` + function baz() { + console.log('baz') + `, + 'typescript' + ) + const diagnostic = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(1, 21, 1, 22), + "'}' expected." + ), + ] + const position = new vscode.Position(1, 22) + + await testDiagnostics( + testDocument, + diagnostic, + position, + 1, + ` + "function baz() { + console.log('baz') + ^ '}' expected." + ` + ) + }) + + it('should only display context within the context lines window for a big file', async () => { + const bigFileContent = Array(100).fill('// Some code here').join('\n') + const testDocument = document( + bigFileContent + + '\n' + + dedent` + function largeFunction() { + let x: number = 5; + let y: string = 'hello'; + let z = x + y; + console.log(x); + } + `, + 'typescript' + ) + const diagnostic = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(103, 16, 103, 21), + "The '+' operator cannot be applied to types 'number' and 'string'." + ), + ] + const position = new vscode.Position(101, 8) + + const { message } = await testDiagnostics( + testDocument, + diagnostic, + position, + 1, + ` + "function largeFunction() { + let x: number = 5; + let y: string = 'hello'; + let z = x + y; + ^^^ The '+' operator cannot be applied to types 'number' and 'string'. + console.log(x); + }" + ` + ) + // Ensure that only the relevant context is shown + expect(message.diagnostic.message).not.toContain('// Some code here') + }) + + it('should handle diagnostics with multiple related information', async () => { + const testDocument = document( + dedent` + function foo(x: number) { + return x.toString(); + } + + let y = foo('5'); + let z = foo(true); + `, + 'typescript' + ) + const diagnostics = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(4, 12, 4, 15), + "Argument of type 'string' is not assignable to parameter of type 'number'.", + 'ts', + [ + { + location: new vscode.Location(testDocument.uri, new vscode.Range(0, 13, 0, 19)), + message: "The expected type comes from parameter 'x' which is declared here", + }, + { + location: new vscode.Location(testDocument.uri, new vscode.Range(0, 13, 0, 19)), + message: "Parameter 'x' is declared as type 'number'", + }, + ] + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(4, 12, 4, 16), + "Argument of type 'boolean' is not assignable to parameter of type 'number'.", + 'ts', + [ + { + location: new vscode.Location(testDocument.uri, new vscode.Range(0, 13, 0, 19)), + message: "The function 'foo' expects a number as its argument", + }, + ] + ), + ] + const position = new vscode.Position(4, 12) + + const { message } = await testDiagnostics( + testDocument, + diagnostics, + position, + 1, + ` + "return x.toString(); + } + + let y = foo('5'); + ^^^ Argument of type 'string' is not assignable to parameter of type 'number'. + ^^^^ Argument of type 'boolean' is not assignable to parameter of type 'number'. + let z = foo(true);" + ` + ) + const relatedErrorList = parser.parse(message.diagnostic.related_information_list) + expect(relatedErrorList[0].message).toContain( + "The expected type comes from parameter 'x' which is declared here" + ) + expect(relatedErrorList[1].message).toContain("Parameter 'x' is declared as type 'number'") + expect(relatedErrorList[2].message).toContain( + "The function 'foo' expects a number as its argument" + ) + }) + it('should return snippets sorted by absolute distance from the current position', async () => { + const testDocument = document( + dedent` + function foo() { + console.log('foo') + } + + function bar() { + let x: number = 'string'; + } + + function baz() { + let y: boolean = 42; + } + + function qux() { + let z: string = true; + } + `, + 'typescript' + ) + const diagnostics = [ + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(5, 24, 5, 32), + "Type 'string' is not assignable to type 'number'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(9, 24, 9, 26), + "Type 'number' is not assignable to type 'boolean'." + ), + createDiagnostic( + vscode.DiagnosticSeverity.Error, + new vscode.Range(13, 24, 13, 28), + "Type 'boolean' is not assignable to type 'string'." + ), + ] + const position = new vscode.Position(10, 0) + + const snippets = await retriever.getDiagnosticsPromptFromInformation( + testDocument, + position, + diagnostics + ) + expect(snippets).toHaveLength(3) + expect(snippets[0].startLine).toBe(9) + expect(snippets[1].startLine).toBe(13) + expect(snippets[2].startLine).toBe(5) + }) +}) diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts new file mode 100644 index 00000000000..978357effca --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts @@ -0,0 +1,174 @@ +import type { AutocompleteContextSnippet } from '@sourcegraph/cody-shared' +import { XMLBuilder } from 'fast-xml-parser' +import * as vscode from 'vscode' +import type { ContextRetriever, ContextRetrieverOptions } from '../../../types' +import { RetrieverIdentifier } from '../../utils' + +// XML builder instance for formatting diagnostic messages +const XML_BUILDER = new XMLBuilder({ format: true }) +// Number of lines of context to include around the diagnostic information in the prompt +const CONTEXT_LINES = 3 + +interface DiagnosticInfo { + message: string + line: number + relatedInformation?: vscode.DiagnosticRelatedInformation[] +} + +export class DiagnosticsRetriever implements vscode.Disposable, ContextRetriever { + public identifier = RetrieverIdentifier.DiagnosticsRetriever + private disposables: vscode.Disposable[] = [] + + public async retrieve({ + document, + position, + }: ContextRetrieverOptions): Promise { + const diagnostics = vscode.languages.getDiagnostics(document.uri) + return this.getDiagnosticsPromptFromInformation(document, position, diagnostics) + } + + public async getDiagnosticsPromptFromInformation( + document: vscode.TextDocument, + position: vscode.Position, + diagnostics: vscode.Diagnostic[] + ): Promise { + const relevantDiagnostics = diagnostics.filter( + diagnostic => diagnostic.severity === vscode.DiagnosticSeverity.Error + ) + const diagnosticInfos = this.getDiagnosticInfos(document, relevantDiagnostics).sort( + (a, b) => Math.abs(a.line - position.line) - Math.abs(b.line - position.line) + ) + return Promise.all( + diagnosticInfos.map(async info => ({ + identifier: this.identifier, + content: await this.getDiagnosticPromptMessage(info), + uri: document.uri, + startLine: info.line, + endLine: info.line, + })) + ) + } + + private getDiagnosticInfos( + document: vscode.TextDocument, + diagnostics: vscode.Diagnostic[] + ): DiagnosticInfo[] { + const diagnosticsByLine = this.getDiagnosticsByLine(diagnostics) + const diagnosticInfos: DiagnosticInfo[] = [] + + for (const [line, lineDiagnostics] of diagnosticsByLine) { + const diagnosticText = this.getDiagnosticsText(document, lineDiagnostics) + if (diagnosticText) { + diagnosticInfos.push({ + message: diagnosticText, + line, + relatedInformation: lineDiagnostics.flatMap(d => d.relatedInformation || []), + }) + } + } + + return diagnosticInfos + } + + private getDiagnosticsByLine(diagnostics: vscode.Diagnostic[]): Map { + const map = new Map() + for (const diagnostic of diagnostics) { + const line = diagnostic.range.start.line + if (!map.has(line)) { + map.set(line, []) + } + map.get(line)!.push(diagnostic) + } + return map + } + + private async getDiagnosticPromptMessage(info: DiagnosticInfo): Promise { + const xmlObj: Record = { + message: info.message, + related_information_list: info.relatedInformation + ? await this.getRelatedInformationPrompt(info.relatedInformation) + : undefined, + } + return XML_BUILDER.build({ diagnostic: xmlObj }) + } + + private async getRelatedInformationPrompt( + relatedInformation: vscode.DiagnosticRelatedInformation[] + ): Promise { + const relatedInfoList = await Promise.all( + relatedInformation.map(async info => { + const document = await vscode.workspace.openTextDocument(info.location.uri) + return { + message: info.message, + file: info.location.uri.fsPath, + text: document.getText(info.location.range), + } + }) + ) + return XML_BUILDER.build(relatedInfoList) + } + + private getDiagnosticsText( + document: vscode.TextDocument, + diagnostics: vscode.Diagnostic[] + ): string | undefined { + if (diagnostics.length === 0) { + return undefined + } + const diagnosticTextList = diagnostics.map(d => this.getDiagnosticMessage(document, d)) + const diagnosticText = diagnosticTextList.join('\n') + const diagnosticLine = diagnostics[0].range.start.line + + return this.addSurroundingContext(document, diagnosticLine, diagnosticText) + } + + private addSurroundingContext( + document: vscode.TextDocument, + diagnosticLine: number, + diagnosticText: string + ): string { + const contextStartLine = Math.max(0, diagnosticLine - CONTEXT_LINES) + const contextEndLine = Math.min(document.lineCount - 1, diagnosticLine + CONTEXT_LINES) + const prevLines = document.getText( + new vscode.Range( + contextStartLine, + 0, + diagnosticLine, + document.lineAt(diagnosticLine).range.end.character + ) + ) + const nextLines = document.getText( + new vscode.Range( + diagnosticLine + 1, + 0, + contextEndLine, + document.lineAt(contextEndLine).range.end.character + ) + ) + return `${prevLines}\n${diagnosticText}\n${nextLines}` + } + + private getDiagnosticMessage(document: vscode.TextDocument, diagnostic: vscode.Diagnostic): string { + const line = document.lineAt(diagnostic.range.start.line) + const column = Math.max(0, diagnostic.range.start.character - 1) + const diagnosticLength = Math.max( + 1, + Math.min( + document.offsetAt(diagnostic.range.end) - document.offsetAt(diagnostic.range.start), + line.text.length + 1 - column + ) + ) + return `${' '.repeat(column)}${'^'.repeat(diagnosticLength)} ${diagnostic.message}` + } + + public isSupportedForLanguageId(): boolean { + return true + } + + public dispose(): void { + for (const disposable of this.disposables) { + disposable.dispose() + } + this.disposables = [] + } +} diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts new file mode 100644 index 00000000000..dccd87354ca --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts @@ -0,0 +1,160 @@ +import dedent from 'dedent' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type * as vscode from 'vscode' +import { Position, Selection } from '../../../../testutils/mocks' +import { document } from '../../../test-helpers' +import { RecentCopyRetriever } from './recent-copy' + +const FIVE_MINUTES = 5 * 60 * 1000 +const MAX_SELECTIONS = 2 + +const disposable = { + dispose: () => {}, +} + +describe('RecentCopyRetriever', () => { + let retriever: RecentCopyRetriever + let onDidChangeTextEditorSelection: any + let mockClipboardContent: string + + const createMockSelection = ( + startLine: number, + startChar: number, + endLine: number, + endChar: number + ) => new Selection(new Position(startLine, startChar), new Position(endLine, endChar)) + + const createMockSelectionForDocument = (document: vscode.TextDocument) => { + return createMockSelection( + 0, + 0, + document.lineCount - 1, + document.lineAt(document.lineCount - 1).text.length + ) + } + + const getDocumentWithUri = (content: string, uri: string, language = 'typescript') => { + return document(content, language, uri) + } + + const simulateSelectionChange = async (testDocument: vscode.TextDocument, selection: Selection) => { + await onDidChangeTextEditorSelection({ + textEditor: { document: testDocument }, + selections: [selection], + }) + // Preloading is debounced so we need to advance the timer manually + await vi.advanceTimersToNextTimerAsync() + } + + beforeEach(() => { + vi.useFakeTimers() + + retriever = new RecentCopyRetriever( + { + maxAgeMs: FIVE_MINUTES, + maxSelections: MAX_SELECTIONS, + }, + { + // Mock VS Code event handlers so we can fire them manually + onDidChangeTextEditorSelection: (_onDidChangeTextEditorSelection: any) => { + onDidChangeTextEditorSelection = _onDidChangeTextEditorSelection + return disposable + }, + } + ) + // Mock the getClipboardContent method to get the vscode clipboard content + vi.spyOn(retriever, 'getClipboardContent').mockImplementation(() => + Promise.resolve(mockClipboardContent) + ) + }) + + afterEach(() => { + retriever.dispose() + }) + + it('should retrieve the copied text if it exists in tracked selections', async () => { + const testDocument = document(dedent` + function foo() { + console.log('foo') + } + `) + mockClipboardContent = testDocument.getText() + const selection = createMockSelectionForDocument(testDocument) + await simulateSelectionChange(testDocument, selection) + const snippets = await retriever.retrieve() + + expect(snippets).toHaveLength(1) + expect(snippets[0]).toEqual({ + content: mockClipboardContent, + uri: testDocument.uri, + startLine: selection.start.line, + endLine: selection.end.line, + identifier: retriever.identifier, + }) + }) + + it('should return null when copied content is not in tracked selections', async () => { + const doc1 = getDocumentWithUri('document 1 content', 'doc1.ts') + const doc2 = getDocumentWithUri('document 2 content', 'doc2.ts') + const doc3 = getDocumentWithUri('document 3 content', 'doc3.ts') + + await simulateSelectionChange(doc1, createMockSelectionForDocument(doc1)) + await simulateSelectionChange(doc2, createMockSelectionForDocument(doc2)) + await simulateSelectionChange(doc3, createMockSelectionForDocument(doc3)) + + mockClipboardContent = doc1.getText() + const snippets = await retriever.retrieve() + + expect(snippets).toHaveLength(0) + }) + + it('should respect maxAgeMs and remove old selections', async () => { + const doc1 = getDocumentWithUri('old content', 'doc1.ts') + await simulateSelectionChange(doc1, createMockSelectionForDocument(doc1)) + vi.advanceTimersByTime(FIVE_MINUTES + 1000) // Advance time beyond maxAgeMs + const doc2 = getDocumentWithUri('new content', 'doc2.ts') + await simulateSelectionChange(doc2, createMockSelectionForDocument(doc2)) + + const trackedSelections = retriever.getTrackedSelections() + expect(trackedSelections).toHaveLength(1) + expect(trackedSelections[0].content).toBe('new content') + }) + + it('should keep tracked selections sorted by timestamp', async () => { + const doc1 = getDocumentWithUri('document 1 content', 'doc1.ts') + const doc2 = getDocumentWithUri('document 2 content', 'doc2.ts') + const doc3 = getDocumentWithUri('document 3 content', 'doc3.ts') + + await simulateSelectionChange(doc1, createMockSelectionForDocument(doc1)) + await simulateSelectionChange(doc2, createMockSelectionForDocument(doc2)) + await simulateSelectionChange(doc3, createMockSelectionForDocument(doc3)) + + const trackedSelections = retriever.getTrackedSelections() + + expect(trackedSelections).toHaveLength(2) + expect(trackedSelections[0].content).toBe('document 3 content') + expect(trackedSelections[1].content).toBe('document 2 content') + }) + + it('should remove outdated selections when scrolling through a document', async () => { + const doc = document(dedent` + line1 + line2 + line3 + line4 + line5 + `) + + // Simulate scrolling through the document + for (let i = 0; i < 5; i++) { + const selection = createMockSelection(0, 0, i, 5) // Select each line + await simulateSelectionChange(doc, selection) + } + + const trackedSelections = retriever.getTrackedSelections() + + // We expect only the most recent selections to be kept (default is 2) + expect(trackedSelections).toHaveLength(1) + expect(trackedSelections[0].content).toBe(doc.getText()) + }) +}) diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts new file mode 100644 index 00000000000..da6c0891160 --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts @@ -0,0 +1,127 @@ +import type { AutocompleteContextSnippet } from '@sourcegraph/cody-shared' +import { debounce } from 'lodash' +import * as vscode from 'vscode' +import type { ContextRetriever } from '../../../types' +import { RetrieverIdentifier } from '../../utils' + +interface TrackedSelection { + timestamp: number + content: string + languageId: string + uri: vscode.Uri + startPosition: vscode.Position + endPosition: vscode.Position +} + +interface RecentCopyRetrieverOptions { + maxAgeMs: number + maxSelections: number +} + +export class RecentCopyRetriever implements vscode.Disposable, ContextRetriever { + public identifier = RetrieverIdentifier.RecentCopyRetriever + private disposables: vscode.Disposable[] = [] + private trackedSelections: TrackedSelection[] = [] + + private readonly maxAgeMs: number + private readonly maxSelections: number + + constructor( + options: RecentCopyRetrieverOptions, + private window: Pick = vscode.window + ) { + this.maxAgeMs = options.maxAgeMs + this.maxSelections = options.maxSelections + + const onSelectionChange = debounce(this.onDidChangeTextEditorSelection.bind(this), 500) + this.disposables.push(this.window.onDidChangeTextEditorSelection(onSelectionChange)) + } + + public async retrieve(): Promise { + const clipboardContent = await this.getClipboardContent() + const selectionItem = this.getSelectionItemIfExist(clipboardContent) + if (selectionItem) { + const autocompleteItem: AutocompleteContextSnippet = { + identifier: this.identifier, + content: selectionItem.content, + uri: selectionItem.uri, + startLine: selectionItem.startPosition.line, + endLine: selectionItem.endPosition.line, + } + return [autocompleteItem] + } + return [] + } + + // This is seperate function because we mock the function in tests + public async getClipboardContent(): Promise { + return vscode.env.clipboard.readText() + } + + public getTrackedSelections(): TrackedSelection[] { + return this.trackedSelections + } + + public isSupportedForLanguageId(): boolean { + return true + } + + private getSelectionItemIfExist(text: string): TrackedSelection | undefined { + return this.trackedSelections.find(ts => ts.content === text) + } + + private addSelectionForTracking(document: vscode.TextDocument, selection: vscode.Selection): void { + if (selection.isEmpty) { + return + } + const selectedText = document.getText(selection) + + const newSelection: TrackedSelection = { + timestamp: Date.now(), + content: selectedText, + languageId: document.languageId, + uri: document.uri, + startPosition: selection.start, + endPosition: selection.end, + } + + this.updateTrackedSelections(newSelection) + } + + private updateTrackedSelections(newSelection: TrackedSelection): void { + const now = Date.now() + this.trackedSelections = this.trackedSelections.filter( + selection => + now - selection.timestamp < this.maxAgeMs && !this.isOverlapping(selection, newSelection) + ) + + this.trackedSelections.unshift(newSelection) + this.trackedSelections = this.trackedSelections.slice(0, this.maxSelections) + } + + // Even with debounce, there is a chance that the same selection is added multiple times if user is slowly selecting + // In that case, we should remove the older selections + private isOverlapping(selection: TrackedSelection, newSelection: TrackedSelection): boolean { + if (selection.uri.toString() !== newSelection.uri.toString()) { + return false + } + return ( + newSelection.startPosition.isBeforeOrEqual(selection.startPosition) && + newSelection.endPosition.isAfterOrEqual(selection.endPosition) + ) + } + + private onDidChangeTextEditorSelection(event: vscode.TextEditorSelectionChangeEvent): void { + const editor = event.textEditor + const selection = event.selections[0] + this.addSelectionForTracking(editor.document, selection) + } + + public dispose(): void { + this.trackedSelections = [] + for (const disposable of this.disposables) { + disposable.dispose() + } + this.disposables = [] + } +} diff --git a/vscode/src/completions/context/retrievers/recent-edits/recent-edits-retriever.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.test.ts similarity index 100% rename from vscode/src/completions/context/retrievers/recent-edits/recent-edits-retriever.test.ts rename to vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.test.ts diff --git a/vscode/src/completions/context/retrievers/recent-edits/recent-edits-retriever.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts similarity index 99% rename from vscode/src/completions/context/retrievers/recent-edits/recent-edits-retriever.ts rename to vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts index 09262ebb6af..6237d511810 100644 --- a/vscode/src/completions/context/retrievers/recent-edits/recent-edits-retriever.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts @@ -48,7 +48,7 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever const content = diff.diff.toString() const autocompleteSnippet = { uri: diff.uri, - identifier: RetrieverIdentifier.RecentEditsRetriever, + identifier: this.identifier, content, } satisfies Omit autocompleteContextSnippets.push(autocompleteSnippet) diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts new file mode 100644 index 00000000000..70dac0e55e0 --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts @@ -0,0 +1,179 @@ +import dedent from 'dedent' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import * as vscode from 'vscode' +import { Position, Range } from '../../../../testutils/mocks' +import { getCurrentDocContext } from '../../../get-current-doc-context' +import { document } from '../../../test-helpers' +import type { ContextRetrieverOptions } from '../../../types' +import { RecentViewPortRetriever } from './recent-view-port' + +const MAX_TRACKED_FILES = 2 + +const documentList = [ + document( + dedent` + function hello() { + console.log('Hello, world!'); + } + `, + 'typescript', + 'file:///test1.ts' + ), + document( + dedent` + class TestClass { + constructor() { + this.name = 'Test'; + } + } + `, + 'typescript', + 'file:///test2.ts' + ), + document( + dedent` + const numbers = [1, 2, 3, 4, 5]; + const sum = numbers.reduce((a, b) => a + b, 0); + `, + 'typescript', + 'file:///test3.ts' + ), +] + +describe('RecentViewPortRetriever', () => { + let retriever: RecentViewPortRetriever + let onDidChangeTextEditorVisibleRanges: any + + const createMockVisibleRange = (doc: vscode.TextDocument, startLine: number, endLine: number) => { + return new Range( + new Position(startLine, 0), + new Position(endLine, doc.lineAt(endLine).text.length) + ) + } + + const getContextRetrieverOptionsFromDoc = (doc: vscode.TextDocument): ContextRetrieverOptions => { + return { + document: doc, + position: new Position(0, 0), + docContext: getCurrentDocContext({ + document: doc, + position: new Position(0, 0), + maxPrefixLength: 100, + maxSuffixLength: 0, + }), + } + } + + beforeEach(() => { + vi.useFakeTimers() + + vi.spyOn(vscode.workspace, 'openTextDocument').mockImplementation(((uri: vscode.Uri) => { + if (uri?.toString().includes('test1.ts')) { + return Promise.resolve(documentList[0]) + } + if (uri?.toString().includes('test2.ts')) { + return Promise.resolve(documentList[1]) + } + if (uri?.toString().includes('test3.ts')) { + return Promise.resolve(documentList[2]) + } + return Promise.resolve(documentList[0]) + }) as any) + + retriever = new RecentViewPortRetriever(MAX_TRACKED_FILES, { + onDidChangeTextEditorVisibleRanges: (_onDidChangeTextEditorVisibleRanges: any) => { + onDidChangeTextEditorVisibleRanges = _onDidChangeTextEditorVisibleRanges + return { dispose: () => {} } + }, + }) + }) + + afterEach(() => { + retriever.dispose() + }) + + const simulateVisibleRangeChange = async ( + testDocument: vscode.TextDocument, + visibleRanges: vscode.Range[] + ) => { + onDidChangeTextEditorVisibleRanges({ + textEditor: { document: testDocument }, + visibleRanges, + }) + // Preloading is debounced so we need to advance the timer manually + await vi.advanceTimersToNextTimerAsync() + } + + it('should ignore the current document', async () => { + const doc = documentList[1] + const visibleRange = createMockVisibleRange(doc, 1, 2) + await simulateVisibleRangeChange(doc, [visibleRange]) + + const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc)) + + expect(snippets).toHaveLength(0) + }) + + it('should retrieve the most recent visible range', async () => { + const doc = documentList[1] + const visibleRange = createMockVisibleRange(doc, 1, 2) + await simulateVisibleRangeChange(doc, [visibleRange]) + const doc2 = documentList[0] + const visibleRange2 = createMockVisibleRange(doc2, 0, 1) + await simulateVisibleRangeChange(doc2, [visibleRange2]) + + const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc2)) + + expect(snippets).toHaveLength(1) + expect(snippets[0]).toMatchObject({ + uri: doc.uri, + startLine: 1, + endLine: 2, + identifier: retriever.identifier, + }) + expect(snippets[0].content).toMatchInlineSnapshot(dedent` + " constructor() { + this.name = 'Test';" + `) + }) + + it('should update existing viewport when revisited', async () => { + const doc = documentList[0] + await simulateVisibleRangeChange(doc, [createMockVisibleRange(doc, 0, 1)]) + await simulateVisibleRangeChange(doc, [createMockVisibleRange(doc, 1, 2)]) + const doc2 = documentList[1] + const visibleRange2 = createMockVisibleRange(doc2, 0, 1) + await simulateVisibleRangeChange(doc2, [visibleRange2]) + + const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc2)) + + expect(snippets).toHaveLength(1) + expect(snippets[0].startLine).toBe(1) + expect(snippets[0].endLine).toBe(2) + }) + + it('should handle empty visible ranges', async () => { + const doc = documentList[0] + await simulateVisibleRangeChange(doc, []) + + const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc)) + + expect(snippets).toHaveLength(0) + }) + + it('should respect MAX_TRACKED_FILES limit', async () => { + const doc1 = documentList[0] + const doc2 = documentList[1] + const doc3 = documentList[2] + + await simulateVisibleRangeChange(doc1, [createMockVisibleRange(doc1, 0, 1)]) + await simulateVisibleRangeChange(doc2, [createMockVisibleRange(doc2, 0, 1)]) + await simulateVisibleRangeChange(doc3, [createMockVisibleRange(doc3, 0, 1)]) + + const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc1)) + + expect(snippets).toHaveLength(2) + expect(snippets[0].uri).toEqual(doc3.uri) + expect(snippets[1].uri).toEqual(doc2.uri) + }) +}) diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts new file mode 100644 index 00000000000..2e42ed63bd9 --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts @@ -0,0 +1,113 @@ +import type { AutocompleteContextSnippet } from '@sourcegraph/cody-shared' +import debounce from 'lodash/debounce' +import { LRUCache } from 'lru-cache' +import * as vscode from 'vscode' +import type { ContextRetriever, ContextRetrieverOptions } from '../../../types' +import { RetrieverIdentifier, type ShouldUseContextParams, shouldBeUsedAsContext } from '../../utils' + +const MAX_RETRIEVED_VIEWPORTS = 5 + +interface TrackedViewPort { + uri: vscode.Uri + visibleRange: vscode.Range + languageId: string + lastAccessTimestamp: number +} + +export class RecentViewPortRetriever implements vscode.Disposable, ContextRetriever { + public identifier = RetrieverIdentifier.RecentViewPortRetriever + private disposables: vscode.Disposable[] = [] + private viewportsByDocumentUri: LRUCache + + constructor( + private readonly maxTrackedFiles: number = 10, + private window: Pick = vscode.window + ) { + this.viewportsByDocumentUri = new LRUCache({ + max: this.maxTrackedFiles, + }) + this.disposables.push( + this.window.onDidChangeTextEditorVisibleRanges( + debounce(this.onDidChangeTextEditorVisibleRanges.bind(this), 300) + ) + ) + } + + public async retrieve({ document }: ContextRetrieverOptions): Promise { + const sortedViewPorts = this.getValidViewPorts(document) + + const snippetPromises = sortedViewPorts.map(async viewPort => { + const document = await vscode.workspace.openTextDocument(viewPort.uri) + const content = document.getText(viewPort.visibleRange) + + return { + uri: viewPort.uri, + content, + startLine: viewPort.visibleRange.start.line, + endLine: viewPort.visibleRange.end.line, + identifier: this.identifier, + } + }) + return Promise.all(snippetPromises) + } + private getValidViewPorts(document: vscode.TextDocument): TrackedViewPort[] { + const currentFileUri = document.uri.toString() + const currentLanguageId = document.languageId + const viewPorts = Array.from(this.viewportsByDocumentUri.entries()) + .map(([_, value]) => value) + .filter((value): value is TrackedViewPort => value !== undefined) + + const sortedViewPorts = viewPorts + .filter(viewport => viewport.uri.toString() !== currentFileUri) + .filter(viewport => { + const params: ShouldUseContextParams = { + enableExtendedLanguagePool: false, + baseLanguageId: currentLanguageId, + languageId: viewport.languageId, + } + return shouldBeUsedAsContext(params) + }) + .sort((a, b) => b.lastAccessTimestamp - a.lastAccessTimestamp) + .slice(0, MAX_RETRIEVED_VIEWPORTS) + + return sortedViewPorts + } + public isSupportedForLanguageId(): boolean { + return true + } + + private onDidChangeTextEditorVisibleRanges(event: vscode.TextEditorVisibleRangesChangeEvent): void { + const { textEditor, visibleRanges } = event + if (visibleRanges.length === 0) { + return + } + const uri = textEditor.document.uri + const visibleRange = visibleRanges[0] + const languageId = textEditor.document.languageId + this.updateTrackedViewPort(uri, visibleRange, languageId) + } + + private updateTrackedViewPort( + uri: vscode.Uri, + visibleRange: vscode.Range, + languageId: string + ): void { + const now = Date.now() + const key = uri.toString() + + this.viewportsByDocumentUri.set(key, { + uri, + visibleRange, + languageId, + lastAccessTimestamp: now, + }) + } + + public dispose(): void { + this.viewportsByDocumentUri.clear() + for (const disposable of this.disposables) { + disposable.dispose() + } + this.disposables = [] + } +} diff --git a/vscode/src/completions/context/utils.ts b/vscode/src/completions/context/utils.ts index ce7a73b98e9..34f0a4a55e2 100644 --- a/vscode/src/completions/context/utils.ts +++ b/vscode/src/completions/context/utils.ts @@ -20,6 +20,9 @@ export enum RetrieverIdentifier { JaccardSimilarityRetriever = 'jaccard-similarity', TscRetriever = 'tsc', LspLightRetriever = 'lsp-light', + RecentCopyRetriever = 'recent-copy', + DiagnosticsRetriever = 'diagnostics', + RecentViewPortRetriever = 'recent-view-port', } export interface ShouldUseContextParams { diff --git a/vscode/src/supercompletions/get-supercompletion.ts b/vscode/src/supercompletions/get-supercompletion.ts index 56222b06a0a..2999f147856 100644 --- a/vscode/src/supercompletions/get-supercompletion.ts +++ b/vscode/src/supercompletions/get-supercompletion.ts @@ -10,7 +10,7 @@ import { import levenshtein from 'js-levenshtein' import * as uuid from 'uuid' import * as vscode from 'vscode' -import type { RecentEditsRetriever } from '../completions/context/retrievers/recent-edits/recent-edits-retriever' +import type { RecentEditsRetriever } from '../completions/context/retrievers/recent-user-actions/recent-edits-retriever' import { ASSISTANT_EXAMPLE, HUMAN_EXAMPLE, MODEL, PROMPT, SYSTEM } from './prompt' import { fixIndentation } from './utils/fix-indentation' import { fuzzyFindLocation } from './utils/fuzzy-find-location' diff --git a/vscode/src/supercompletions/supercompletion-provider.ts b/vscode/src/supercompletions/supercompletion-provider.ts index 98ebdda813d..fe8b06f19d1 100644 --- a/vscode/src/supercompletions/supercompletion-provider.ts +++ b/vscode/src/supercompletions/supercompletion-provider.ts @@ -1,6 +1,6 @@ import type { ChatClient } from '@sourcegraph/cody-shared' import * as vscode from 'vscode' -import { RecentEditsRetriever } from '../completions/context/retrievers/recent-edits/recent-edits-retriever' +import { RecentEditsRetriever } from '../completions/context/retrievers/recent-user-actions/recent-edits-retriever' import type { CodyStatusBar } from '../services/StatusBar' import { type Supercompletion, getSupercompletions } from './get-supercompletion' import { SupercompletionRenderer } from './renderer'