fix recent edits context source (#6071)

## Context 1. Recent edits context source outputs wrong diffs sometimes. The behaviour is not reproducible always, but can be seen from the output logs. Adding one such example in below: <img width="1442" alt="image" src="https://github.com/user-attachments/assets/d36c8601-3a8e-42f4-a5e8-c3276a6e9e35"> 2. On further inspection, it turns out that sometimes when events are received `private onDidChangeTextDocument` the text document could have changed, and reading `document.getText()` returns change with the updated text (i.e. including the change we received). Although this behaviour is not always reproducible. So, this PR reads the text document during the constructor call or when the document is first opened. ## Test plan Manual inspection
sourcegraph · Nov 8, 2024 · 1167681 · 1167681
1 parent 07df3f9
commit 1167681
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 7 deletions.
diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.test.ts
@@ -15,6 +15,7 @@ describe('RecentEditsRetriever', () => {
     let onDidChangeTextDocument: (event: vscode.TextDocumentChangeEvent) => void
     let onDidRenameFiles: (event: vscode.FileRenameEvent) => void
     let onDidDeleteFiles: (event: vscode.FileDeleteEvent) => void
+    let onDidOpenTextDocument: (event: vscode.TextDocument) => void
 
     beforeEach(() => {
         vi.useFakeTimers()
@@ -38,6 +39,10 @@ describe('RecentEditsRetriever', () => {
                     onDidDeleteFiles = listener
                     return { dispose: () => {} }
                 },
+                onDidOpenTextDocument(listener) {
+                    onDidOpenTextDocument = listener
+                    return { dispose: () => {} }
+                },
             }
         )
     })
@@ -100,6 +105,8 @@ describe('RecentEditsRetriever', () => {
         `)
 
         it('tracks document changes and creates a git diff', async () => {
+            onDidOpenTextDocument(testDocument)
+
             replaceFooLogWithNumber(testDocument)
 
             deleteBarLog(testDocument)
@@ -125,6 +132,9 @@ describe('RecentEditsRetriever', () => {
 
         it('no-ops for blocked files due to the context filter', async () => {
             vi.spyOn(contextFiltersProvider, 'isUriIgnored').mockResolvedValueOnce('repo:foo')
+
+            onDidOpenTextDocument(testDocument)
+
             replaceFooLogWithNumber(testDocument)
 
             deleteBarLog(testDocument)
@@ -135,6 +145,8 @@ describe('RecentEditsRetriever', () => {
         })
 
         it('does not yield changes that are older than the configured timeout', async () => {
+            onDidOpenTextDocument(testDocument)
+
             replaceFooLogWithNumber(testDocument)
 
             vi.advanceTimersByTime(3 * 60 * 1000)
@@ -159,6 +171,8 @@ describe('RecentEditsRetriever', () => {
         })
 
         it('handles renames', async () => {
+            onDidOpenTextDocument(testDocument)
+
             replaceFooLogWithNumber(testDocument)
 
             vi.advanceTimersByTime(3 * 60 * 1000)
@@ -194,8 +208,12 @@ describe('RecentEditsRetriever', () => {
         })
 
         it('handles deletions', async () => {
+            onDidOpenTextDocument(testDocument)
+
             replaceFooLogWithNumber(testDocument)
+
             onDidDeleteFiles({ files: [testDocument.uri] })
+
             expect(await retriever.getDiff(testDocument.uri)).toBe(null)
         })
     })

diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts
@@ -38,14 +38,20 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever
         options: RecentEditsRetrieverOptions,
         readonly workspace: Pick<
             typeof vscode.workspace,
-            'onDidChangeTextDocument' | 'onDidRenameFiles' | 'onDidDeleteFiles'
+            'onDidChangeTextDocument' | 'onDidRenameFiles' | 'onDidDeleteFiles' | 'onDidOpenTextDocument'
         > = vscode.workspace
     ) {
         this.maxAgeMs = options.maxAgeMs
         this.addLineNumbersForDiff = options.addLineNumbersForDiff ?? false
+        // Track the already open documents when editor was opened
+        for (const document of vscode.workspace.textDocuments) {
+            this.trackDocument(document)
+        }
+        // Set up event listeners for changes
         this.disposables.push(workspace.onDidChangeTextDocument(this.onDidChangeTextDocument.bind(this)))
         this.disposables.push(workspace.onDidRenameFiles(this.onDidRenameFiles.bind(this)))
         this.disposables.push(workspace.onDidDeleteFiles(this.onDidDeleteFiles.bind(this)))
+        this.disposables.push(workspace.onDidOpenTextDocument(this.onDidOpenTextDocument.bind(this)))
     }
 
     public async retrieve(options: ContextRetrieverOptions): Promise<AutocompleteContextSnippet[]> {
@@ -76,7 +82,7 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever
         const diffPromises = Array.from(this.trackedDocuments.entries()).map(
             async ([uri, trackedDocument]) => {
                 const diff = await this.getDiff(vscode.Uri.parse(uri))
-                if (diff) {
+                if (diff && trackedDocument.changes.length > 0) {
                     return {
                         diff,
                         uri: trackedDocument.uri,
@@ -177,9 +183,9 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever
     }
 
     private onDidChangeTextDocument(event: vscode.TextDocumentChangeEvent): void {
-        let trackedDocument = this.trackedDocuments.get(event.document.uri.toString())
+        const trackedDocument = this.trackedDocuments.get(event.document.uri.toString())
         if (!trackedDocument) {
-            trackedDocument = this.trackDocument(event.document)
+            return
         }
 
         const now = Date.now()
@@ -216,15 +222,17 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever
         }
     }
 
-    private trackDocument(document: vscode.TextDocument): TrackedDocument {
+    private trackDocument(document: vscode.TextDocument): void {
+        if (document.uri.scheme !== 'file') {
+            return
+        }
         const trackedDocument: TrackedDocument = {
             content: document.getText(),
             languageId: document.languageId,
             uri: document.uri,
             changes: [],
         }
         this.trackedDocuments.set(document.uri.toString(), trackedDocument)
-        return trackedDocument
     }
 
     private reconcileOutdatedChanges(): void {
@@ -242,6 +250,12 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever
             trackedDocument.changes = trackedDocument.changes.slice(firstNonOutdatedChangeIndex)
         }
     }
+
+    private onDidOpenTextDocument(document: vscode.TextDocument): void {
+        if (!this.trackedDocuments.has(document.uri.toString())) {
+            this.trackDocument(document)
+        }
+    }
 }
 
 function applyChanges(content: string, changes: vscode.TextDocumentContentChangeEvent[]): string {

diff --git a/vscode/src/supercompletions/supercompletion-provider.ts b/vscode/src/supercompletions/supercompletion-provider.ts
@@ -23,7 +23,7 @@ export class SupercompletionProvider implements vscode.Disposable {
         },
         readonly workspace: Pick<
             typeof vscode.workspace,
-            'onDidChangeTextDocument' | 'onDidRenameFiles' | 'onDidDeleteFiles'
+            'onDidChangeTextDocument' | 'onDidRenameFiles' | 'onDidDeleteFiles' | 'onDidOpenTextDocument'
         > = vscode.workspace
     ) {
         this.renderer = new SupercompletionRenderer()