fix: pass right URIs from withConnection()

Author: felixfbecker

src/connection.ts

@@ -7,6 +7,7 @@ import {
     WebSocketMessageReader,
     WebSocketMessageWriter,
 } from '@sourcegraph/vscode-ws-jsonrpc'
+import { attempt } from 'lodash'
 import { fromEvent, merge, Subject } from 'rxjs'
 import { filter, map, mapTo, take } from 'rxjs/operators'
 import { Subscribable, Unsubscribable } from 'sourcegraph'
@@ -70,8 +71,8 @@ export const webSocketTransport = ({
                 map(({ params }) => params)
             ),
         unsubscribe: () => {
-            socket.close()
-            connection.dispose()
+            attempt(() => socket.close())
+            attempt(() => connection.dispose())
         },
     }
 }

src/features/feature.ts

@@ -23,18 +23,19 @@ export type RegistrationOptions<T extends RequestType<any, any, any, any>> = Exc
 
 export function scopeDocumentSelectorToRoot(
     documentSelector: DocumentSelector | null,
-    rootUri: URL | null
+    clientRootUri: URL | null
 ): DocumentSelector {
     if (!documentSelector || documentSelector.length === 0) {
         documentSelector = [{ pattern: '**' }]
     }
-    if (!rootUri) {
+    if (!clientRootUri) {
         return documentSelector
     }
     return documentSelector
         .map((filter): DocumentFilter => (typeof filter === 'string' ? { language: filter } : filter))
         .map(filter => ({
             ...filter,
-            pattern: new URL(filter.pattern || '**', rootUri).href,
+            // TODO filter.pattern needs to be run resolved relative to server root URI before mounting on clientRootUri
+            pattern: new URL(filter.pattern || '**', clientRootUri).href,
         }))
 }

src/index.ts

@@ -60,7 +60,7 @@ export interface LSPClient extends Unsubscribable {
      * Ensures a connection with the given workspace root, passes it to the given function.
      * If the workspace is not currently open in Sourcegraph, the connection is closed again after the Promise returned by the function resolved.
      *
-     * @param workspaceRoot The workspace folder root URI that will be ensured to be open before calling the function.
+     * @param workspaceRoot The client workspace folder root URI that will be ensured to be open before calling the function.
      * @param fn Callback that is called with the connection.
      */
     withConnection<R>(workspaceRoot: URL, fn: (connection: LSPConnection) => Promise<R>): Promise<R>
@@ -128,7 +128,7 @@ export async function register({
 
     const registrationSubscriptions = new Map<string, Unsubscribable>()
     /**
-     * @param scopeRootUri A workspace folder root URI to scope the providers to. If `null`, the provider is registered for all workspace folders.
+     * @param scopeRootUri A client workspace folder root URI to scope the providers to. If `null`, the provider is registered for all workspace folders.
      */
     function registerCapabilities(
         connection: LSPConnection,
@@ -153,7 +153,7 @@ export async function register({
         }
     }
 
-    async function connect(): Promise<LSPConnection> {
+    async function connect(clientRootUri: URL | null, initParams: InitializeParams): Promise<LSPConnection> {
         const subscriptions = new Subscription()
         const decorationType = sourcegraph.app.createDecorationType()
         const connection = await createConnection()
@@ -256,12 +256,13 @@ export async function register({
                     }
                 })
         )
+        await initializeConnection(connection, clientRootUri, initParams)
         return connection
     }
 
     async function initializeConnection(
         connection: LSPConnection,
-        rootUri: URL | null,
+        clientRootUri: URL | null,
         initParams: InitializeParams
     ): Promise<void> {
         const initializeResult = await connection.sendRequest(InitializeRequest.type, initParams)
@@ -273,18 +274,26 @@ export async function register({
 
         // Listen for dynamic capabilities
         connection.setRequestHandler(RegistrationRequest.type, params => {
-            registerCapabilities(connection, rootUri, params.registrations)
+            registerCapabilities(connection, clientRootUri, params.registrations)
         })
         // Register static capabilities
-        registerCapabilities(connection, rootUri, staticRegistrations)
+        registerCapabilities(connection, clientRootUri, staticRegistrations)
 
         await afterInitialize(initializeResult)
     }
 
     let withConnection: <R>(workspaceFolder: URL, fn: (connection: LSPConnection) => Promise<R>) => Promise<R>
 
     if (supportsWorkspaceFolders) {
-        const connection = await connect()
+        const connection = await connect(
+            null,
+            {
+                processId: null,
+                rootUri: null,
+                capabilities: clientCapabilities,
+                workspaceFolders: sourcegraph.workspace.roots.map(toLSPWorkspaceFolder({ clientToServerURI })),
+            }
+        )
         subscriptions.add(connection)
         withConnection = async (workspaceFolder, fn) => {
             let tempWorkspaceFolder: WorkspaceFolder | undefined
@@ -310,12 +319,6 @@ export async function register({
                 }
             }
         }
-        await initializeConnection(connection, null, {
-            processId: null,
-            rootUri: null,
-            capabilities: clientCapabilities,
-            workspaceFolders: sourcegraph.workspace.roots.map(toLSPWorkspaceFolder),
-        })
 
         // Forward root changes
         subscriptions.add(
@@ -328,8 +331,12 @@ export async function register({
                         after,
                     })),
                     map(({ before, after }) => ({
-                        added: differenceBy(after, before, root => root.uri.toString()).map(toLSPWorkspaceFolder),
-                        removed: differenceBy(before, after, root => root.uri.toString()).map(toLSPWorkspaceFolder),
+                        added: differenceBy(after, before, root => root.uri.toString()).map(
+                            toLSPWorkspaceFolder({ clientToServerURI })
+                        ),
+                        removed: differenceBy(before, after, root => root.uri.toString()).map(
+                            toLSPWorkspaceFolder({ clientToServerURI })
+                        ),
                     }))
                 )
                 .subscribe(event => {
@@ -338,31 +345,47 @@ export async function register({
         )
     } else {
         // Supports only one workspace root
+        // TODO this should store a refcount to avoid closing connections other consumers have a reference to
+        /** Map from client root URI to connection */
         const connectionsByRootUri = new Map<string, Promise<LSPConnection>>()
         withConnection = async (workspaceFolder, fn) => {
             let connection = await connectionsByRootUri.get(workspaceFolder.href)
-            if (!connection) {
-                connection = await connect()
-                subscriptions.add(connection)
+            if (connection) {
+                return await fn(connection)
             }
+            const serverRootUri = clientToServerURI(workspaceFolder)
+            connection = await connect(
+                workspaceFolder,
+                {
+                    processId: null,
+                    rootUri: serverRootUri.href,
+                    capabilities: clientCapabilities,
+                    workspaceFolders: null,
+                }
+            )
+            subscriptions.add(connection)
             try {
                 return await fn(connection)
             } finally {
+                connectionsByRootUri.delete(workspaceFolder.href)
                 connection.unsubscribe()
             }
         }
         function addRoots(added: ReadonlyArray<WorkspaceRoot>): void {
             for (const root of added) {
                 const connectionPromise = (async () => {
                     try {
-                        const connection = await connect()
+                        const serverRootUri = clientToServerURI(new URL(root.uri.toString()))
+                        const connection = await connect(
+                            new URL(root.uri.toString()),
+                            {
+                                processId: null,
+                                rootUri: serverRootUri.href,
+                                capabilities: clientCapabilities,
+                                workspaceFolders: null,
+                            }
+                        )
                         subscriptions.add(connection)
-                        await initializeConnection(connection, new URL(root.uri.toString()), {
-                            processId: null,
-                            rootUri: root.uri.toString(),
-                            capabilities: clientCapabilities,
-                            workspaceFolders: null,
-                        })
                         return connection
                     } catch (err) {
                         logger.error('Error creating connection', err)

src/lsp-conversion.ts

@@ -94,10 +94,15 @@ export const convertDiagnosticToDecoration = (
     range: convertRange(sourcegraph, diagnostic.range),
 })
 
-export const toLSPWorkspaceFolder = (root: sourcegraph.WorkspaceRoot): WorkspaceFolder => ({
-    uri: root.uri.toString(),
-    name: new URL(root.uri.toString()).pathname.split('/').pop()!,
-})
+export const toLSPWorkspaceFolder = ({ clientToServerURI }: { clientToServerURI: (u: URL) => URL }) => (
+    root: sourcegraph.WorkspaceRoot
+): WorkspaceFolder => {
+    const serverUri = clientToServerURI(new URL(root.uri.toString()))
+    return {
+        uri: serverUri.href,
+        name: new URL(serverUri.href).pathname.split('/').pop()!,
+    }
+}
 
 /**
  * Rewrites all `uri` properties in an object, recursively