feat: initial support for local project context (#949)

* feat: initial support for local project context

* feat: dynamically import vector library

* fix: various cleanup

* fix: path test errors on Windows

* feat: add support for enablement via LSP configuration

* fix: better qualified name for local project enablement configuration

Author: breedloj

.gitignore

@@ -5,6 +5,7 @@ build
 **/*.tgz
 !core/codewhisperer-streaming/amzn-codewhisperer-streaming-*.tgz
 !core/q-developer-streaming-client/amzn-amazon-q-developer-streaming-client-*.tgz
+!server/aws-lsp-codewhisperer/types/types-local-indexing-*.tgz
 
 .testresults/**
 

app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts

@@ -6,6 +6,7 @@ import {
     QChatServerProxy,
     QConfigurationServerTokenProxy,
     QNetTransformServerTokenProxy,
+    QLocalProjectContextServerTokenProxy,
 } from '@aws/lsp-codewhisperer'
 import { IdentityServer } from '@aws/lsp-identity'
 
@@ -23,6 +24,7 @@ const props: RuntimeProps = {
         QNetTransformServerTokenProxy,
         QChatServerProxy,
         IdentityServer.create,
+        QLocalProjectContextServerTokenProxy,
     ],
     name: 'AWS CodeWhisperer',
 }

package-lock.json

@@ -51,6 +51,7 @@
     "devDependencies": {
         "@types/adm-zip": "^0.5.5",
         "@types/archiver": "^6.0.2",
+        "@types/local-indexing": "file:./types/types-local-indexing-1.0.0.tgz",
         "@types/lokijs": "^1.5.14",
         "@types/uuid": "^9.0.8",
         "@types/diff": "^7.0.2",

server/aws-lsp-codewhisperer/package.json

@@ -0,0 +1,217 @@
+import { LocalProjectContextController } from './localProjectContextController'
+import { SinonStub, stub, spy, assert as sinonAssert, match } from 'sinon'
+import * as assert from 'assert'
+import * as fs from 'fs'
+import { Dirent } from 'fs'
+import * as path from 'path'
+import { URI } from 'vscode-uri'
+
+class LoggingMock {
+    public error: SinonStub
+    public info: SinonStub
+    public log: SinonStub
+    public warn: SinonStub
+
+    constructor() {
+        this.error = stub()
+        this.info = stub()
+        this.log = stub()
+        this.warn = stub()
+    }
+}
+
+describe('LocalProjectContextController', () => {
+    let controller: LocalProjectContextController
+    let logging: LoggingMock
+    let mockWorkspaceFolders: any[]
+    let vectorLibMock: any
+    let fsStub: SinonStub
+
+    const BASE_PATH = path.join(__dirname, 'path', 'to', 'workspace1')
+
+    beforeEach(() => {
+        logging = new LoggingMock()
+        mockWorkspaceFolders = [
+            {
+                uri: URI.file(BASE_PATH),
+                name: 'workspace1',
+            },
+        ]
+
+        vectorLibMock = {
+            start: stub().resolves({
+                buildIndex: stub().resolves(),
+                clear: stub().resolves(),
+                queryVectorIndex: stub().resolves(['mockChunk1', 'mockChunk2']),
+                queryInlineProjectContext: stub().resolves(['mockContext1']),
+                updateIndexV2: stub().resolves(),
+            }),
+        }
+
+        fsStub = stub(fs.promises, 'readdir')
+        fsStub.withArgs(match.string, { withFileTypes: true }).callsFake(path => {
+            if (path.endsWith('workspace1')) {
+                return Promise.resolve([createMockDirent('Test.java', false), createMockDirent('src', true)])
+            } else if (path.endsWith('src')) {
+                return Promise.resolve([createMockDirent('Main.java', false)])
+            } else {
+                return Promise.resolve([])
+            }
+        })
+        controller = new LocalProjectContextController('testClient', mockWorkspaceFolders, logging as any)
+    })
+
+    afterEach(() => {
+        fsStub.restore()
+    })
+
+    describe('init', () => {
+        it('should initialize vector library successfully', async () => {
+            await controller.init(vectorLibMock)
+
+            sinonAssert.notCalled(logging.error)
+            sinonAssert.called(vectorLibMock.start)
+            const vecLib = await vectorLibMock.start()
+            sinonAssert.called(vecLib.buildIndex)
+        })
+
+        it('should handle initialization errors', async () => {
+            vectorLibMock.start.rejects(new Error('Init failed'))
+
+            await controller.init(vectorLibMock)
+
+            sinonAssert.called(logging.error)
+        })
+    })
+
+    describe('queryVectorIndex', () => {
+        beforeEach(async () => {
+            await controller.init(vectorLibMock)
+        })
+
+        it('should return empty array when vector library is not initialized', async () => {
+            const uninitializedController = new LocalProjectContextController(
+                'testClient',
+                mockWorkspaceFolders,
+                logging as any
+            )
+
+            const result = await uninitializedController.queryVectorIndex({ query: 'test' })
+            assert.deepStrictEqual(result, [])
+        })
+
+        it('should return chunks from vector library', async () => {
+            const result = await controller.queryVectorIndex({ query: 'test' })
+            assert.deepStrictEqual(result, ['mockChunk1', 'mockChunk2'])
+        })
+
+        it('should handle query errors', async () => {
+            const vecLib = await vectorLibMock.start()
+            vecLib.queryVectorIndex.rejects(new Error('Query failed'))
+
+            const result = await controller.queryVectorIndex({ query: 'test' })
+            assert.deepStrictEqual(result, [])
+            sinonAssert.called(logging.error)
+        })
+    })
+
+    describe('queryInlineProjectContext', () => {
+        beforeEach(async () => {
+            await controller.init(vectorLibMock)
+        })
+
+        it('should return empty array when vector library is not initialized', async () => {
+            const uninitializedController = new LocalProjectContextController(
+                'testClient',
+                mockWorkspaceFolders,
+                logging as any
+            )
+
+            const result = await uninitializedController.queryInlineProjectContext({
+                query: 'test',
+                filePath: 'test.java',
+                target: 'test',
+            })
+            assert.deepStrictEqual(result, [])
+        })
+
+        it('should return context from vector library', async () => {
+            const result = await controller.queryInlineProjectContext({
+                query: 'test',
+                filePath: 'test.java',
+                target: 'test',
+            })
+            assert.deepStrictEqual(result, ['mockContext1'])
+        })
+
+        it('should handle query errors', async () => {
+            const vecLib = await vectorLibMock.start()
+            vecLib.queryInlineProjectContext.rejects(new Error('Query failed'))
+
+            const result = await controller.queryInlineProjectContext({
+                query: 'test',
+                filePath: 'test.java',
+                target: 'test',
+            })
+            assert.deepStrictEqual(result, [])
+            sinonAssert.called(logging.error)
+        })
+    })
+
+    describe('updateIndex', () => {
+        beforeEach(async () => {
+            await controller.init(vectorLibMock)
+        })
+
+        it('should do nothing when vector library is not initialized', async () => {
+            const uninitializedController = new LocalProjectContextController(
+                'testClient',
+                mockWorkspaceFolders,
+                logging as any
+            )
+
+            await uninitializedController.updateIndex(['test.java'], 'add')
+            sinonAssert.notCalled(logging.error)
+        })
+
+        it('should update index successfully', async () => {
+            const vecLib = await vectorLibMock.start()
+            await controller.updateIndex(['test.java'], 'add')
+            sinonAssert.called(vecLib.updateIndexV2)
+        })
+
+        it('should handle update errors', async () => {
+            const vecLib = await vectorLibMock.start()
+            vecLib.updateIndexV2.rejects(new Error('Update failed'))
+
+            await controller.updateIndex(['test.java'], 'add')
+            sinonAssert.called(logging.error)
+        })
+    })
+
+    describe('dispose', () => {
+        it('should clear and remove vector library reference', async () => {
+            await controller.init(vectorLibMock)
+            await controller.dispose()
+
+            const vecLib = await vectorLibMock.start()
+            sinonAssert.called(vecLib.clear)
+
+            const queryResult = await controller.queryVectorIndex({ query: 'test' })
+            assert.deepStrictEqual(queryResult, [])
+        })
+    })
+})
+
+function createMockDirent(name: string, isDirectory: boolean): Dirent {
+    return {
+        name,
+        isDirectory: () => isDirectory,
+        isFile: () => !isDirectory,
+        isBlockDevice: () => false,
+        isCharacterDevice: () => false,
+        isFIFO: () => false,
+        isSocket: () => false,
+        isSymbolicLink: () => false,
+    } as Dirent
+}