Send initial YAML and JSON files right after language client starts (2/3)

extension/src/lspClient/languageClient.ts

@@ -1,12 +1,14 @@
-import { workspace } from 'vscode'
+import { Uri, workspace } from 'vscode'
 import {
   LanguageClient,
   LanguageClientOptions,
   ServerOptions,
   TransportKind
 } from 'vscode-languageclient/node'
 import { documentSelector, serverModule } from 'dvc-vscode-lsp'
+import { readFileSync } from 'fs-extra'
 import { Disposable } from '../class/dispose'
+import { findFiles } from '../fileSystem/workspace'
 
 export class LanguageClientWrapper extends Disposable {
   private client: LanguageClient
@@ -37,8 +39,27 @@ export class LanguageClientWrapper extends Disposable {
     this.start()
   }
 
-  start() {
-    this.client.start()
+  async start() {
+    await this.client.start()
+
+    const files = await findFiles('**/*.{yaml,json,py,toml}', '.??*')
+
+    const textDocuments = files.map(filePath => {
+      const uri = Uri.file(filePath).toString()
+      const languageId = filePath.endsWith('yaml') ? 'yaml' : 'json'
+      const text = readFileSync(filePath, 'utf8')
+
+      return {
+        languageId,
+        text,
+        uri,
+        version: 0
+      }
+    })
+
+    await this.client.sendRequest('initialTextDocuments', {
+      textDocuments
+    })
 
     return this
   }

languageServer/src/LanguageServer.ts

@@ -10,17 +10,19 @@ import {
   Location,
   Position,
   Range,
-  DocumentSymbol
+  DocumentSymbol,
+  TextDocumentItem
 } from 'vscode-languageserver/node'
 import { TextDocument } from 'vscode-languageserver-textdocument'
 import { URI } from 'vscode-uri'
 import { TextDocumentWrapper } from './TextDocumentWrapper'
 
 export class LanguageServer {
-  private documents!: TextDocuments<TextDocument>
+  private documentsKnownToEditor!: TextDocuments<TextDocument>
+  private documentsFromDvcClient: TextDocumentWrapper[] = []
 
   public listen(connection: _Connection) {
-    this.documents = new TextDocuments(TextDocument)
+    this.documentsKnownToEditor = new TextDocuments(TextDocument)
 
     connection.onInitialize(() => this.onInitialize())
 
@@ -32,20 +34,59 @@ export class LanguageServer {
       return this.onDefinition(params)
     })
 
-    this.documents.listen(connection)
+    connection.onRequest(
+      'initialTextDocuments',
+      (params: { textDocuments: TextDocumentItem[] }) => {
+        this.documentsFromDvcClient = params.textDocuments.map(
+          ({ uri, languageId, version, text: content }) => {
+            const textDocument = TextDocument.create(
+              uri,
+              languageId,
+              version,
+              content
+            )
+
+            return this.wrap(textDocument)
+          }
+        )
+      }
+    )
+
+    this.documentsKnownToEditor.listen(connection)
 
     connection.listen()
   }
 
+  private getAllDocuments() {
+    const openDocuments = this.documentsKnownToEditor.all()
+    const acc: TextDocumentWrapper[] = openDocuments.map(doc => this.wrap(doc))
+
+    for (const textDocument of this.documentsFromDvcClient) {
+      const userAlreadyOpenedIt = this.documentsKnownToEditor.get(
+        textDocument.uri
+      )
+
+      if (!userAlreadyOpenedIt) {
+        acc.push(textDocument)
+      }
+    }
+
+    return acc
+  }
+
   private getDvcTextDocument(
     params: TextDocumentPositionParams | CodeActionParams
   ) {
     const uri = params.textDocument.uri
-    const doc = this.documents.get(uri)
+    const doc = this.documentsKnownToEditor.get(uri)
 
     if (!doc) {
-      return null
+      const alternative = this.documentsFromDvcClient.find(
+        txtDoc => txtDoc.uri === uri
+      )
+      return alternative ?? null
     }
+
     return this.wrap(doc)
   }
 
@@ -55,7 +96,7 @@ export class LanguageServer {
 
   private getFilePathLocations(
     symbolUnderCursor: DocumentSymbol,
-    allDocs: TextDocument[]
+    allDocs: TextDocumentWrapper[]
   ) {
     if (symbolUnderCursor.kind !== SymbolKind.File) {
       return []
@@ -79,13 +120,12 @@ export class LanguageServer {
 
   private getLocationsFromOtherDocuments(
     symbolUnderCursor: DocumentSymbol,
-    allDocs: TextDocument[]
+    allDocs: TextDocumentWrapper[]
   ) {
     const locationsAccumulator = []
 
     for (const txtDoc of allDocs) {
-      const finder = this.wrap(txtDoc)
-      const locations = finder.findLocationsFor(symbolUnderCursor)
+      const locations = txtDoc.findLocationsFor(symbolUnderCursor)
       locationsAccumulator.push(...locations)
     }
 
@@ -97,7 +137,7 @@ export class LanguageServer {
     const symbolUnderCursor = document?.symbolAt(params.position)
 
     if (document && symbolUnderCursor) {
-      const allDocs = this.documents.all()
+      const allDocs = this.getAllDocuments()
       const locationsAccumulator = []
 
       const fileLocations = this.getFilePathLocations(

languageServer/src/test/definitions.test.ts

@@ -11,6 +11,7 @@ import {
   disposeTestConnections,
   setupTestConnections
 } from './utils/setup-test-connections'
+import { sendTheseFilesToServer } from './utils/sendTheseFilesToServer'
 
 describe('textDocument/definitions', () => {
   beforeEach(() => {
@@ -75,7 +76,7 @@ describe('textDocument/definitions', () => {
   })
 
   it('should provide a single location that points to the top of the file path symbol', async () => {
-    const [dvcYaml] = await openTheseFilesAndNotifyServer([
+    const [dvcYaml] = await sendTheseFilesToServer([
       {
         languageId: 'yaml',
         mockContents: vars_dvc_yaml,

languageServer/src/test/utils/sendTheseFilesToServer.ts

@@ -0,0 +1,27 @@
+import { TextDocument } from 'vscode-languageserver-textdocument'
+import { URI } from 'vscode-uri'
+import { client } from './setup-test-connections'
+
+export const sendTheseFilesToServer = async (
+  files: Array<{ mockPath: string; mockContents: string; languageId: string }>
+) => {
+  const filesToReturn: TextDocument[] = []
+
+  for (const { mockPath, mockContents, languageId } of files) {
+    const uri = URI.file(mockPath).toString()
+    const textDocument = TextDocument.create(uri, languageId, 1, mockContents)
+    filesToReturn.push(textDocument)
+  }
+  await client.sendRequest('initialTextDocuments', {
+    textDocuments: filesToReturn.map(textDocument => {
+      return {
+        languageId: 'yaml',
+        text: textDocument.getText(),
+        uri: textDocument.uri,
+        version: textDocument.version
+      }
+    })
+  })
+
+  return filesToReturn
+}