feat(amazonq): basic paginated suggestions using language server

packages/amazonq/src/app/inline/completion.ts

@@ -13,81 +13,175 @@ import {
     TextDocument,
     commands,
     languages,
+    Disposable,
 } from 'vscode'
 import { LanguageClient } from 'vscode-languageclient'
-import {
-    InlineCompletionListWithReferences,
-    InlineCompletionWithReferencesParams,
-    inlineCompletionWithReferencesRequestType,
-    logInlineCompletionSessionResultsNotificationType,
-    LogInlineCompletionSessionResultsParams,
-} from '@aws/language-server-runtimes/protocol'
+import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-runtimes/protocol'
+import { SessionManager } from './sessionManager'
+import { RecommendationService } from './recommendationService'
 import { CodeWhispererConstants } from 'aws-core-vscode/codewhisperer'
 
-export function registerInlineCompletion(languageClient: LanguageClient) {
-    const inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(languageClient)
-    languages.registerInlineCompletionItemProvider(CodeWhispererConstants.platformLanguageIds, inlineCompletionProvider)
+export class InlineCompletionManager implements Disposable {
+    private disposable: Disposable
+    private inlineCompletionProvider: AmazonQInlineCompletionItemProvider
+    private languageClient: LanguageClient
+    private sessionManager: SessionManager
+    private recommendationService: RecommendationService
+    private readonly logSessionResultMessageName = 'aws/logInlineCompletionSessionResults'
+
+    constructor(languageClient: LanguageClient) {
+        this.languageClient = languageClient
+        this.sessionManager = new SessionManager()
+        this.recommendationService = new RecommendationService(this.sessionManager)
+        this.inlineCompletionProvider = new AmazonQInlineCompletionItemProvider(
+            languageClient,
+            this.recommendationService,
+            this.sessionManager
+        )
+        this.disposable = languages.registerInlineCompletionItemProvider(
+            CodeWhispererConstants.platformLanguageIds,
+            this.inlineCompletionProvider
+        )
+    }
+
+    public dispose(): void {
+        if (this.disposable) {
+            this.disposable.dispose()
+        }
+    }
+
+    public registerInlineCompletion() {
+        const onInlineAcceptance = async (
+            sessionId: string,
+            itemId: string,
+            requestStartTime: number,
+            firstCompletionDisplayLatency?: number
+        ) => {
+            // TODO: also log the seen state for other suggestions in session
+            const params: LogInlineCompletionSessionResultsParams = {
+                sessionId: sessionId,
+                completionSessionResult: {
+                    [itemId]: {
+                        seen: true,
+                        accepted: true,
+                        discarded: false,
+                    },
+                },
+                totalSessionDisplayTime: Date.now() - requestStartTime,
+                firstCompletionDisplayLatency: firstCompletionDisplayLatency,
+            }
+            this.languageClient.sendNotification(this.logSessionResultMessageName, params)
+            this.disposable.dispose()
+            this.disposable = languages.registerInlineCompletionItemProvider(
+                CodeWhispererConstants.platformLanguageIds,
+                this.inlineCompletionProvider
+            )
+        }
+        commands.registerCommand('aws.amazonq.acceptInline', onInlineAcceptance)
 
-    const onInlineAcceptance = async (
-        sessionId: string,
-        itemId: string,
-        requestStartTime: number,
-        firstCompletionDisplayLatency?: number
-    ) => {
-        const params: LogInlineCompletionSessionResultsParams = {
-            sessionId: sessionId,
-            completionSessionResult: {
-                [itemId]: {
-                    seen: true,
-                    accepted: true,
-                    discarded: false,
+        const onInlineRejection = async () => {
+            await commands.executeCommand('editor.action.inlineSuggest.hide')
+            // TODO: also log the seen state for other suggestions in session
+            this.disposable.dispose()
+            this.disposable = languages.registerInlineCompletionItemProvider(
+                CodeWhispererConstants.platformLanguageIds,
+                this.inlineCompletionProvider
+            )
+            const sessionId = this.sessionManager.getActiveSession()?.sessionId
+            const itemId = this.sessionManager.getActiveRecommendation()[0]?.itemId
+            if (!sessionId || !itemId) {
+                return
+            }
+            const params: LogInlineCompletionSessionResultsParams = {
+                sessionId: sessionId,
+                completionSessionResult: {
+                    [itemId]: {
+                        seen: true,
+                        accepted: false,
+                        discarded: false,
+                    },
                 },
-            },
-            totalSessionDisplayTime: Date.now() - requestStartTime,
-            firstCompletionDisplayLatency: firstCompletionDisplayLatency,
+            }
+            this.languageClient.sendNotification(this.logSessionResultMessageName, params)
         }
-        languageClient.sendNotification(logInlineCompletionSessionResultsNotificationType as any, params)
+        commands.registerCommand('aws.amazonq.rejectCodeSuggestion', onInlineRejection)
+
+        /*
+            We have to overwrite the prev. and next. commands because the inlineCompletionProvider only contained the current item
+            To show prev. and next. recommendation we need to re-register a new provider with the previous or next item
+        */
+
+        const swapProviderAndShow = async () => {
+            await commands.executeCommand('editor.action.inlineSuggest.hide')
+            this.disposable.dispose()
+            this.disposable = languages.registerInlineCompletionItemProvider(
+                CodeWhispererConstants.platformLanguageIds,
+                new AmazonQInlineCompletionItemProvider(
+                    this.languageClient,
+                    this.recommendationService,
+                    this.sessionManager,
+                    false
+                )
+            )
+            await commands.executeCommand('editor.action.inlineSuggest.trigger')
+        }
+
+        const prevCommandHandler = async () => {
+            this.sessionManager.decrementActiveIndex()
+            await swapProviderAndShow()
+        }
+        commands.registerCommand('editor.action.inlineSuggest.showPrevious', prevCommandHandler)
+
+        const nextCommandHandler = async () => {
+            this.sessionManager.incrementActiveIndex()
+            await swapProviderAndShow()
+        }
+        commands.registerCommand('editor.action.inlineSuggest.showNext', nextCommandHandler)
     }
-    commands.registerCommand('aws.sample-vscode-ext-amazonq.accept', onInlineAcceptance)
 }
 
 export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
-    constructor(private readonly languageClient: LanguageClient) {}
+    constructor(
+        private readonly languageClient: LanguageClient,
+        private readonly recommendationService: RecommendationService,
+        private readonly sessionManager: SessionManager,
+        private readonly isNewSession: boolean = true
+    ) {}
 
     async provideInlineCompletionItems(
         document: TextDocument,
         position: Position,
         context: InlineCompletionContext,
         token: CancellationToken
     ): Promise<InlineCompletionItem[] | InlineCompletionList> {
-        const requestStartTime = Date.now()
-        const request: InlineCompletionWithReferencesParams = {
-            textDocument: {
-                uri: document.uri.toString(),
-            },
-            position,
-            context,
+        if (this.isNewSession) {
+            // make service requests if it's a new session
+            await this.recommendationService.getAllRecommendations(
+                this.languageClient,
+                document,
+                position,
+                context,
+                token
+            )
         }
-
-        const response = await this.languageClient.sendRequest(
-            inlineCompletionWithReferencesRequestType as any,
-            request,
-            token
-        )
-
-        const list: InlineCompletionListWithReferences = response as InlineCompletionListWithReferences
-        this.languageClient.info(`Client: Received ${list.items.length} suggestions`)
-        const firstCompletionDisplayLatency = Date.now() - requestStartTime
-
-        // Add completion session tracking and attach onAcceptance command to each item to record used decision
-        for (const item of list.items) {
+        // get active item from session for displaying
+        const items = this.sessionManager.getActiveRecommendation()
+        const session = this.sessionManager.getActiveSession()
+        if (!session || !items.length) {
+            return []
+        }
+        for (const item of items) {
             item.command = {
-                command: 'aws.sample-vscode-ext-amazonq.accept',
+                command: 'aws.amazonq.acceptInline',
                 title: 'On acceptance',
-                arguments: [list.sessionId, item.itemId, requestStartTime, firstCompletionDisplayLatency],
+                arguments: [
+                    session.sessionId,
+                    item.itemId,
+                    session.requestStartTime,
+                    session.firstCompletionDisplayLatency,
+                ],
             }
         }
-
-        return list as InlineCompletionList
+        return items as InlineCompletionItem[]
     }
 }

packages/amazonq/src/app/inline/recommendationService.ts

@@ -0,0 +1,78 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+    InlineCompletionListWithReferences,
+    InlineCompletionWithReferencesParams,
+    inlineCompletionWithReferencesRequestType,
+} from '@aws/language-server-runtimes/protocol'
+import { CancellationToken, InlineCompletionContext, Position, TextDocument } from 'vscode'
+import { LanguageClient } from 'vscode-languageclient'
+import { SessionManager } from './sessionManager'
+
+export class RecommendationService {
+    constructor(private readonly sessionManager: SessionManager) {}
+
+    async getAllRecommendations(
+        languageClient: LanguageClient,
+        document: TextDocument,
+        position: Position,
+        context: InlineCompletionContext,
+        token: CancellationToken
+    ) {
+        const request: InlineCompletionWithReferencesParams = {
+            textDocument: {
+                uri: document.uri.toString(),
+            },
+            position,
+            context,
+        }
+        const requestStartTime = Date.now()
+
+        // Handle first request
+        const firstResult: InlineCompletionListWithReferences = await languageClient.sendRequest(
+            inlineCompletionWithReferencesRequestType as any,
+            request,
+            token
+        )
+
+        const firstCompletionDisplayLatency = Date.now() - requestStartTime
+        this.sessionManager.startSession(
+            firstResult.sessionId,
+            firstResult.items,
+            requestStartTime,
+            firstCompletionDisplayLatency
+        )
+
+        if (firstResult.partialResultToken) {
+            // If there are more results to fetch, handle them in the background
+            this.processRemainingRequests(languageClient, request, firstResult, token).catch((error) => {
+                languageClient.warn(`Error when getting suggestions: ${error}`)
+            })
+        } else {
+            this.sessionManager.closeSession()
+        }
+    }
+
+    private async processRemainingRequests(
+        languageClient: LanguageClient,
+        initialRequest: InlineCompletionWithReferencesParams,
+        firstResult: InlineCompletionListWithReferences,
+        token: CancellationToken
+    ): Promise<void> {
+        let nextToken = firstResult.partialResultToken
+        while (nextToken) {
+            const request = { ...initialRequest, partialResultToken: nextToken }
+            const result: InlineCompletionListWithReferences = await languageClient.sendRequest(
+                inlineCompletionWithReferencesRequestType as any,
+                request,
+                token
+            )
+            this.sessionManager.updateSessionSuggestions(result.items)
+            nextToken = result.partialResultToken
+        }
+        this.sessionManager.closeSession()
+    }
+}