fix: add support for determing workspace folder with root uri/path on initialize (#1210)

LSP lifecycle for VS extension is controlled by Visual Studio. It currently does not set the workspaceFolder as part of the Initialize Params of the LSP initialize request, however it does send this information via rootUri and rootPath properties. The workspaceFolders is utilized for the local project context indexing and so in order to support that, this change extends logic to determine the workspaceFolder to use the root properties when workspaceFolders is not set.

Author: shruti0085

server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts

@@ -4,6 +4,7 @@ import { TelemetryService } from '../../shared/telemetry/telemetryService'
 import { LocalProjectContextController } from '../../shared/localProjectContextController'
 import { languageByExtension } from '../../shared/languageDetection'
 import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
+import { getWorkspaceFolders } from '../../shared/initializeUtils'
 import { URI } from 'vscode-uri'
 
 export const LocalProjectContextServer =
@@ -16,9 +17,10 @@ export const LocalProjectContextServer =
         let localProjectContextEnabled: boolean = false
 
         lsp.addInitializer((params: InitializeParams) => {
+            const workspaceFolders = getWorkspaceFolders(logging, params)
             localProjectContextController = new LocalProjectContextController(
                 params.clientInfo?.name ?? 'unknown',
-                params.workspaceFolders ?? [],
+                workspaceFolders,
                 logging
             )
 

server/aws-lsp-codewhisperer/src/shared/initializeUtils.test.ts

@@ -0,0 +1,125 @@
+import { InitializeParams, WorkspaceFolder } from '@aws/language-server-runtimes/protocol'
+import { URI } from 'vscode-uri'
+import * as path from 'path'
+import assert = require('assert')
+import sinon = require('sinon')
+import { TestFeatures } from '@aws/language-server-runtimes/testing'
+import { Features } from '@aws/language-server-runtimes/server-interface/server'
+import { getWorkspaceFolders } from './initializeUtils'
+
+describe('initializeUtils', () => {
+    let logging: Features['logging']
+
+    before(function () {
+        logging = new TestFeatures().logging
+    })
+
+    describe('getWorkspaceFolders', () => {
+        const sampleWorkspaceUri = 'file:///path/to/folder'
+        const sampleWorkspaceName = 'folder'
+        const sampleWorkspaceFolder: WorkspaceFolder = {
+            name: sampleWorkspaceName,
+            uri: sampleWorkspaceUri,
+        }
+
+        const createParams = (params: Partial<InitializeParams>) => params as InitializeParams
+
+        it('should return workspaceFolders when provided', () => {
+            const workspaceFolders: WorkspaceFolder[] = [
+                sampleWorkspaceFolder,
+                { name: 'folder2', uri: 'file:///path/to/folder2' },
+            ]
+            const params = createParams({ workspaceFolders })
+            const result = getWorkspaceFolders(logging, params)
+
+            assert.deepStrictEqual(result, workspaceFolders)
+        })
+
+        describe('should create workspace folder from rootUri when workspaceFolders is not provided', () => {
+            const invalidWorkspaceFolderCases = [
+                ['no workspaceFolder param', { rootUri: sampleWorkspaceUri }],
+                ['empty workspaceFolder param params', { WorkspaceFolders: [], rootUri: sampleWorkspaceUri }],
+            ] as const
+
+            invalidWorkspaceFolderCases.forEach(([name, input]) => {
+                it(`should return root uri for ${name}`, () => {
+                    const params = createParams(input)
+                    const result = getWorkspaceFolders(logging, params)
+                    assert.deepStrictEqual(result, [sampleWorkspaceFolder])
+                })
+            })
+            const params = createParams({ rootUri: sampleWorkspaceUri })
+            const result = getWorkspaceFolders(logging, params)
+
+            assert.deepStrictEqual(result, [sampleWorkspaceFolder])
+        })
+
+        it('should create workspace folder from rootPath when neither workspaceFolders nor rootUri is provided', () => {
+            const rootPath = '/path/to/folder'
+            const params = createParams({ rootPath: rootPath })
+            const result = getWorkspaceFolders(logging, params)
+
+            assert.deepStrictEqual(result, [sampleWorkspaceFolder])
+        })
+
+        it('should use uri as folder name when URI basename is empty', () => {
+            const rootUri = 'file:///'
+            const params = createParams({ rootUri })
+            const result = getWorkspaceFolders(logging, params)
+
+            assert.deepStrictEqual(result, [{ name: rootUri, uri: rootUri }])
+        })
+
+        it('should handle Windows paths correctly', () => {
+            const rootPath = 'C:\\Users\\test\\folder'
+            const pathUri = URI.parse(rootPath).toString()
+            const params = createParams({ rootPath })
+
+            const result = getWorkspaceFolders(logging, params)
+
+            const expectedName = path.basename(URI.parse(pathUri).fsPath)
+            assert.deepStrictEqual(result, [{ name: expectedName, uri: pathUri }])
+        })
+
+        it('should handle rootUri with special characters', () => {
+            const rootUri = 'file:///path/to/special%20project'
+            const decodedPath = URI.parse(rootUri).path
+            const folderName = path.basename(decodedPath)
+
+            const params = createParams({ rootUri })
+            const result = getWorkspaceFolders(logging, params)
+
+            assert.deepStrictEqual(result, [{ name: folderName, uri: rootUri }])
+            assert.equal('special project', result[0].name)
+        })
+
+        describe('should return empty workspaceFolder array', () => {
+            const emptyArrayCases = [
+                ['no params', {}],
+                ['undefined params', undefined as unknown as InitializeParams],
+                ['null params', null as unknown as InitializeParams],
+                ['empty workspaceFolders', { workspaceFolders: [] }],
+            ] as const
+
+            emptyArrayCases.forEach(([name, input]) => {
+                it(`should return empty array for ${name}`, () => {
+                    const result = getWorkspaceFolders(logging, input as InitializeParams)
+                    assert.equal(result.length, 0)
+                })
+            })
+        })
+
+        it('should handle errors and return empty array', () => {
+            const basenameStub = sinon.stub(path, 'basename').throws(new Error('Test error'))
+
+            try {
+                const badParams = createParams({ rootUri: sampleWorkspaceUri })
+                const result = getWorkspaceFolders(logging, badParams)
+
+                assert.deepStrictEqual(result, [])
+            } finally {
+                basenameStub.restore()
+            }
+        })
+    })
+})

server/aws-lsp-codewhisperer/src/shared/initializeUtils.ts

@@ -0,0 +1,32 @@
+import { InitializeParams, WorkspaceFolder } from '@aws/language-server-runtimes/protocol'
+import { Logging } from '@aws/language-server-runtimes/server-interface'
+import * as path from 'path'
+import { URI } from 'vscode-uri'
+
+export function getWorkspaceFolders(logging: Logging, params?: InitializeParams): WorkspaceFolder[] {
+    if (!params) {
+        return []
+    }
+
+    if (params.workspaceFolders && params.workspaceFolders.length > 0) {
+        return params.workspaceFolders
+    }
+    try {
+        const getFolderName = (parsedUri: URI) => path.basename(parsedUri.fsPath) || parsedUri.toString()
+
+        if (params.rootUri) {
+            const parsedUri = URI.parse(params.rootUri)
+            const folderName = getFolderName(parsedUri)
+            return [{ name: folderName, uri: params.rootUri }]
+        }
+        if (params.rootPath) {
+            const parsedUri = URI.parse(params.rootPath)
+            const folderName = getFolderName(parsedUri)
+            return [{ name: folderName, uri: parsedUri.toString() }]
+        }
+        return []
+    } catch (error) {
+        logging.error(`Error occurred when determining workspace folders: ${error}`)
+        return []
+    }
+}