Add client side support for refreshing source generated files

Author: dibarbet

l10n/bundle.l10n.json

@@ -163,6 +163,7 @@
   "Fix all issues": "Fix all issues",
   "Select fix all action": "Select fix all action",
   "Test run already in progress": "Test run already in progress",
+  "Generated document not found": "Generated document not found",
   "Server stopped": "Server stopped",
   "Workspace projects": "Workspace projects",
   "Your workspace has multiple Visual Studio Solution files; please select one to get full IntelliSense.": "Your workspace has multiple Visual Studio Solution files; please select one to get full IntelliSense.",

src/lsptoolshost/roslynLanguageServer.ts

@@ -29,6 +29,7 @@ import {
     CancellationToken,
     RequestHandler,
     ResponseError,
+    NotificationHandler0,
 } from 'vscode-languageclient/node';
 import { PlatformInformation } from '../shared/platform';
 import { readConfigurations } from './configurationMiddleware';
@@ -423,6 +424,10 @@ export class RoslynLanguageServer {
         this._languageClient.addDisposable(this._languageClient.onRequest(type, handler));
     }
 
+    public registerOnNotification(method: string, handler: NotificationHandler0) {
+        this._languageClient.addDisposable(this._languageClient.onNotification(method, handler));
+    }
+
     public async registerSolutionSnapshot(token: vscode.CancellationToken): Promise<SolutionSnapshotId> {
         const response = await this.sendRequest0(RoslynProtocol.RegisterSolutionSnapshotRequest.type, token);
         if (response) {

src/lsptoolshost/roslynProtocol.ts

@@ -233,10 +233,12 @@ export interface CopilotRelatedDocumentsReport {
 
 export interface SourceGeneratorGetRequestParams {
     textDocument: lsp.TextDocumentIdentifier;
+    resultId?: string;
 }
 
 export interface SourceGeneratedDocumentText {
-    text: string;
+    text?: string;
+    resultId?: string;
 }
 
 export namespace WorkspaceDebugConfigurationRequest {
@@ -366,3 +368,9 @@ export namespace SourceGeneratorGetTextRequest {
     export const messageDirection: lsp.MessageDirection = lsp.MessageDirection.clientToServer;
     export const type = new lsp.RequestType<SourceGeneratorGetRequestParams, SourceGeneratedDocumentText, void>(method);
 }
+
+export namespace RefreshSourceGeneratedDocumentNotification {
+    export const method = 'workspace/refreshSourceGeneratedDocument';
+    export const messageDirection: lsp.MessageDirection = lsp.MessageDirection.serverToClient;
+    export const type = new lsp.NotificationType(method);
+}

src/lsptoolshost/sourceGeneratedFilesContentProvider.ts

@@ -8,6 +8,7 @@ import * as RoslynProtocol from './roslynProtocol';
 import { RoslynLanguageServer } from './roslynLanguageServer';
 import { UriConverter } from './uriConverter';
 import * as lsp from 'vscode-languageserver-protocol';
+import { IDisposable } from '@microsoft/servicehub-framework';
 
 export function registerSourceGeneratedFilesContentProvider(
     context: vscode.ExtensionContext,
@@ -16,16 +17,91 @@ export function registerSourceGeneratedFilesContentProvider(
     context.subscriptions.push(
         vscode.workspace.registerTextDocumentContentProvider(
             'roslyn-source-generated',
-            new (class implements vscode.TextDocumentContentProvider {
-                async provideTextDocumentContent(uri: vscode.Uri, token: vscode.CancellationToken): Promise<string> {
-                    const result = await languageServer.sendRequest(
-                        RoslynProtocol.SourceGeneratorGetTextRequest.type,
-                        { textDocument: lsp.TextDocumentIdentifier.create(UriConverter.serialize(uri)) },
-                        token
-                    );
-                    return result.text;
-                }
-            })()
+            new RoslynSourceGeneratedContentProvider(languageServer)
         )
     );
 }
+
+class RoslynSourceGeneratedContentProvider implements vscode.TextDocumentContentProvider, IDisposable {
+    private _onDidChangeEmitter: vscode.EventEmitter<vscode.Uri> = new vscode.EventEmitter<vscode.Uri>();
+
+    // Stores all the source generated documents that we have opened so far and their up to date content.
+    private _openedDocuments: Map<vscode.Uri, RoslynProtocol.SourceGeneratedDocumentText> = new Map();
+
+    // Since we could potentially have multiple refresh notifications in flight at the same time,
+    // we use a simple queue to ensure that updates to our state map only happen serially.
+    private _updateQueue?: Promise<RoslynProtocol.SourceGeneratedDocumentText>;
+
+    private _cancellationSource = new vscode.CancellationTokenSource();
+
+    constructor(private languageServer: RoslynLanguageServer) {
+        languageServer.registerOnNotification(
+            RoslynProtocol.RefreshSourceGeneratedDocumentNotification.method,
+            async () => {
+                this._openedDocuments.forEach(async (_, key) => {
+                    await this.enqueueDocumentUpdateAsync(key, this._cancellationSource.token);
+                    this._onDidChangeEmitter.fire(key);
+                });
+            }
+        );
+        vscode.workspace.onDidCloseTextDocument((document) => {
+            const openedDoc = this._openedDocuments.get(document.uri);
+            if (openedDoc !== undefined) {
+                this._openedDocuments.delete(document.uri);
+            }
+        });
+    }
+
+    public onDidChange: vscode.Event<vscode.Uri> = this._onDidChangeEmitter.event;
+
+    async provideTextDocumentContent(uri: vscode.Uri, token: vscode.CancellationToken): Promise<string> {
+        let content = this._openedDocuments.get(uri);
+
+        if (!content) {
+            // We're being asked about this document for the first time, so we need to fetch it from the server.
+            content = await this.enqueueDocumentUpdateAsync(uri, token);
+        }
+
+        return content.text ?? vscode.l10n.t('Generated document not found');
+    }
+
+    private async enqueueDocumentUpdateAsync(
+        uri: vscode.Uri,
+        token: vscode.CancellationToken
+    ): Promise<RoslynProtocol.SourceGeneratedDocumentText> {
+        if (!this._updateQueue) {
+            this._updateQueue = this.updateDocumentAsync(uri, token);
+        } else {
+            this._updateQueue = this._updateQueue.then(async () => await this.updateDocumentAsync(uri, token));
+        }
+
+        return await this._updateQueue;
+    }
+
+    private async updateDocumentAsync(
+        uri: vscode.Uri,
+        token: vscode.CancellationToken
+    ): Promise<RoslynProtocol.SourceGeneratedDocumentText> {
+        const currentContent = this._openedDocuments.get(uri);
+        const newContent = await this.languageServer.sendRequest(
+            RoslynProtocol.SourceGeneratorGetTextRequest.type,
+            {
+                textDocument: lsp.TextDocumentIdentifier.create(UriConverter.serialize(uri)),
+                resultId: currentContent?.resultId,
+            },
+            token
+        );
+
+        // If we had no content before, or the resultId has changed, update the content
+        if (!currentContent || newContent.resultId !== currentContent?.resultId) {
+            this._openedDocuments.set(uri, newContent);
+            return newContent;
+        }
+
+        return currentContent;
+    }
+
+    dispose(): void {
+        this._cancellationSource.cancel();
+    }
+}