feat: merge updates for inline completions (#1299)

Author: zixlin7

client/vscode/src/inlineCompletionActivation.ts

@@ -41,6 +41,12 @@ export const CodewhispererInlineCompletionLanguages = [
     { scheme: 'file', language: 'scala' },
     { scheme: 'file', language: 'vue' },
     { scheme: 'file', language: 'csharp' },
+    { scheme: 'file', language: 'c' },
+    { scheme: 'file', language: 'cpp' },
+    { scheme: 'file', language: 'python' },
+    { scheme: 'file', language: 'sql' },
+    { scheme: 'file', language: 'jsx' },
+    { scheme: 'file', language: 'tsx' },
 ]
 
 export function registerInlineCompletion(languageClient: LanguageClient) {

server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.test.ts

@@ -21,17 +21,22 @@ import {
 import { CodeWhispererSession, SessionData, SessionManager } from './session/sessionManager'
 import {
     EMPTY_RESULT,
+    EXPECTED_NEXT_TOKEN,
     EXPECTED_REFERENCE,
     EXPECTED_RESPONSE_CONTEXT,
     EXPECTED_RESULT,
+    EXPECTED_RESULT_WITHOUT_IMPORTS,
     EXPECTED_RESULT_WITHOUT_REFERENCES,
+    EXPECTED_RESULT_WITH_IMPORTS,
     EXPECTED_RESULT_WITH_REFERENCES,
     EXPECTED_SESSION_ID,
     EXPECTED_SUGGESTION,
     EXPECTED_SUGGESTION_LIST,
+    EXPECTED_SUGGESTION_LIST_WITH_IMPORTS,
     HELLO_WORLD_IN_CSHARP,
     HELLO_WORLD_LINE,
     HELLO_WORLD_WITH_WINDOWS_ENDING,
+    SAMPLE_SESSION_DATA,
     SINGLE_LINE_FILE_CUTOFF_INDEX,
     SOME_CLOSED_FILE,
     SOME_FILE,
@@ -413,8 +418,10 @@ describe('CodeWhisperer Server', () => {
                         insertText: deletedLine.concat('\n    '),
                         range: undefined,
                         references: undefined,
+                        mostRelevantMissingImports: undefined,
                     },
                 ],
+                partialResultToken: undefined,
             }
 
             const result = await features.doInlineCompletionWithReferences(
@@ -497,8 +504,10 @@ describe('CodeWhisperer Server', () => {
                         insertText: HELLO_WORLD_LINE.substring(0, SINGLE_LINE_FILE_CUTOFF_INDEX),
                         range: undefined,
                         references: undefined,
+                        mostRelevantMissingImports: undefined,
                     },
                 ],
+                partialResultToken: undefined,
             }
 
             const result = await features.doInlineCompletionWithReferences(
@@ -529,8 +538,10 @@ describe('CodeWhisperer Server', () => {
                         insertText: EXPECTED_SUGGESTION[0].content,
                         range: undefined,
                         references: undefined,
+                        mostRelevantMissingImports: undefined,
                     },
                 ],
+                partialResultToken: undefined,
             }
 
             const result = await features.doInlineCompletionWithReferences(
@@ -561,6 +572,54 @@ describe('CodeWhisperer Server', () => {
             assert.deepEqual(result, EMPTY_RESULT)
         })
 
+        // pagination
+        it('returns next token from service', async () => {
+            service.generateSuggestions.returns(
+                Promise.resolve({
+                    suggestions: EXPECTED_SUGGESTION,
+                    responseContext: { ...EXPECTED_RESPONSE_CONTEXT, nextToken: EXPECTED_NEXT_TOKEN },
+                })
+            )
+
+            const result = await features.doInlineCompletionWithReferences(
+                {
+                    textDocument: { uri: SOME_FILE.uri },
+                    position: { line: 0, character: 0 },
+                    context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+                },
+                CancellationToken.None
+            )
+
+            assert.deepEqual(result, { ...EXPECTED_RESULT, partialResultToken: EXPECTED_NEXT_TOKEN })
+        })
+
+        it('handles partialResultToken in request', async () => {
+            const manager = SessionManager.getInstance()
+            manager.createSession(SAMPLE_SESSION_DATA)
+            await features.doInlineCompletionWithReferences(
+                {
+                    textDocument: { uri: SOME_FILE.uri },
+                    position: { line: 0, character: 0 },
+                    context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+                    partialResultToken: EXPECTED_NEXT_TOKEN,
+                },
+                CancellationToken.None
+            )
+
+            const expectedGenerateSuggestionsRequest = {
+                fileContext: {
+                    filename: SOME_FILE.uri,
+                    programmingLanguage: { languageName: 'csharp' },
+                    leftFileContent: '',
+                    rightFileContent: HELLO_WORLD_IN_CSHARP,
+                },
+                maxResults: 5,
+                nextToken: EXPECTED_NEXT_TOKEN,
+            }
+
+            sinon.assert.calledOnceWithExactly(service.generateSuggestions, expectedGenerateSuggestionsRequest)
+        })
+
         describe('Supplemental Context', () => {
             it('should send supplemental context when using token authentication', async () => {
                 const test_service = sinon.createStubInstance(
@@ -700,6 +759,53 @@ describe('CodeWhisperer Server', () => {
             assert.deepEqual(result, EXPECTED_RESULT_WITHOUT_REFERENCES)
         })
 
+        it('should not include import statements if no settings are specified', async () => {
+            features.lsp.workspace.getConfiguration.returns(Promise.resolve({}))
+            await startServer(features, server)
+
+            service.generateSuggestions.returns(
+                Promise.resolve({
+                    suggestions: EXPECTED_SUGGESTION_LIST_WITH_IMPORTS,
+                    responseContext: EXPECTED_RESPONSE_CONTEXT,
+                })
+            )
+
+            const result = await features.openDocument(SOME_FILE).doInlineCompletionWithReferences(
+                {
+                    textDocument: { uri: SOME_FILE.uri },
+                    position: { line: 0, character: 0 },
+                    context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+                },
+                CancellationToken.None
+            )
+
+            // Check the completion result
+            assert.deepEqual(result, EXPECTED_RESULT_WITHOUT_IMPORTS)
+        })
+
+        it('should include import statements if enabled', async () => {
+            features.lsp.workspace.getConfiguration.returns(Promise.resolve({ includeImportsWithSuggestions: true }))
+            await startServer(features, server)
+
+            service.generateSuggestions.returns(
+                Promise.resolve({
+                    suggestions: EXPECTED_SUGGESTION_LIST_WITH_IMPORTS,
+                    responseContext: EXPECTED_RESPONSE_CONTEXT,
+                })
+            )
+
+            const result = await features.openDocument(SOME_FILE).doInlineCompletionWithReferences(
+                {
+                    textDocument: { uri: SOME_FILE.uri },
+                    position: { line: 0, character: 0 },
+                    context: { triggerKind: InlineCompletionTriggerKind.Invoked },
+                },
+                CancellationToken.None
+            )
+
+            assert.deepEqual(result, EXPECTED_RESULT_WITH_IMPORTS)
+        })
+
         it('should filter recommendations with references if GetConfiguration is not handled by the client', async () => {
             features.lsp.workspace.getConfiguration.returns(Promise.reject(new Error('GetConfiguration failed')))
             await startServer(features, server)
@@ -870,8 +976,10 @@ describe('CodeWhisperer Server', () => {
                                 },
                             },
                         ],
+                        mostRelevantMissingImports: undefined,
                     },
                 ],
+                partialResultToken: undefined,
             }
             service.generateSuggestions.returns(
                 Promise.resolve({
@@ -929,8 +1037,10 @@ describe('CodeWhisperer Server', () => {
                         insertText: insertText,
                         range: undefined,
                         references: undefined,
+                        mostRelevantMissingImports: undefined,
                     },
                 ],
+                partialResultToken: undefined,
             }
             service.generateSuggestions.returns(
                 Promise.resolve({
@@ -1792,7 +1902,7 @@ describe('CodeWhisperer Server', () => {
             const EXPECTED_COMPLETION_RESPONSES = [
                 { sessionId: '', items: [] },
                 { sessionId: '', items: [] },
-                { sessionId: SESSION_IDS_LOG[2], items: EXPECTED_RESULT.items }, // Last session wins
+                { sessionId: SESSION_IDS_LOG[2], items: EXPECTED_RESULT.items, partialResultToken: undefined }, // Last session wins
             ]
             // Only last request must return completion items
             assert.deepEqual(getCompletionsResponses, EXPECTED_COMPLETION_RESPONSES)