config(inline): show next session if only both sessino ids are the same

packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts

@@ -49,6 +49,7 @@ describe('inlineCompletionService', function () {
             const session = CodeWhispererSessionState.instance.getSession()
             const mockEditor = createMockTextEditor()
             sinon.stub(RecommendationHandler.instance, 'getRecommendations').resolves({
+                sessionId: 'foo',
                 result: 'Succeeded',
                 errorMessage: undefined,
                 recommendationCount: 1,

packages/core/src/codewhisperer/commands/onInlineAcceptance.ts

@@ -152,7 +152,11 @@ async function promoteNextSessionIfAvailable(acceptanceEntry: OnRecommendationAc
         const nextSession = CodeWhispererSessionState.instance.getNextSession()
         nextSession.startPos = acceptanceEntry.editor.selection.active
         CodeWhispererSessionState.instance.setSession(nextSession)
-        if (nextSession.recommendations.length) {
+
+        if (nextSession.recommendations.length && acceptanceEntry.sessionId === nextSession.sessionId) {
+            getLogger().debug(
+                `next session id ${nextSession.sessionId} mismatched previous session id ${acceptanceEntry.sessionId}, aborting promoteNextSession`
+            )
             await RecommendationHandler.instance.tryShowRecommendation()
         }
     }

packages/core/src/codewhisperer/models/model.ts

@@ -81,6 +81,7 @@ export interface CodeWhispererSupplementalContextItem {
 
 // This response struct can contain more info as needed
 export interface GetRecommendationsResponse {
+    readonly sessionId: string
     readonly result: 'Succeeded' | 'Failed'
     readonly recommendationCount: number
     readonly errorMessage: string | undefined

packages/core/src/codewhisperer/service/inlineCompletionService.ts

@@ -85,6 +85,7 @@ export class InlineCompletionService {
     ): Promise<GetRecommendationsResponse> {
         if (vsCodeState.isCodeWhispererEditing || RecommendationHandler.instance.isSuggestionVisible()) {
             return {
+                sessionId: '',
                 result: 'Failed',
                 errorMessage: 'Amazon Q is already running',
                 recommendationCount: 0,
@@ -104,6 +105,7 @@ export class InlineCompletionService {
         if (AuthUtil.instance.isConnectionExpired()) {
             await AuthUtil.instance.notifyReauthenticate(isAutoTrigger)
             return {
+                sessionId: '',
                 result: 'Failed',
                 errorMessage: 'auth',
                 recommendationCount: 0,
@@ -115,6 +117,7 @@ export class InlineCompletionService {
         RecommendationHandler.instance.checkAndResetCancellationTokens()
         RecommendationHandler.instance.documentUri = editor.document.uri
         let response: GetRecommendationsResponse = {
+            sessionId: '',
             result: 'Failed',
             errorMessage: undefined,
             recommendationCount: 0,
@@ -140,6 +143,7 @@ export class InlineCompletionService {
                         void showTimedMessage(response.errorMessage ? response.errorMessage : noSuggestions, 2000)
                     }
                     return {
+                        sessionId: '',
                         result: 'Failed',
                         errorMessage: 'cancelled',
                         recommendationCount: 0,
@@ -160,6 +164,7 @@ export class InlineCompletionService {
         TelemetryHelper.instance.tryRecordClientComponentLatency()
 
         return {
+            sessionId: response.sessionId,
             result: 'Succeeded',
             errorMessage: undefined,
             recommendationCount: session.recommendations.length,

packages/core/src/codewhisperer/service/recommendationHandler.ts

@@ -181,6 +181,7 @@ export class RecommendationHandler {
 
         if (!editor) {
             return Promise.resolve<GetRecommendationsResponse>({
+                sessionId: '',
                 result: invocationResult,
                 errorMessage: errorMessage,
                 recommendationCount: 0,
@@ -259,6 +260,7 @@ export class RecommendationHandler {
                     errorMessage = `${languageName} is currently not supported by Amazon Q inline suggestions`
                 }
                 return Promise.resolve<GetRecommendationsResponse>({
+                    sessionId: '',
                     result: invocationResult,
                     errorMessage: errorMessage,
                     recommendationCount: 0,
@@ -286,6 +288,8 @@ export class RecommendationHandler {
             requestId = resp?.$response && resp?.$response?.requestId
             nextToken = resp?.nextToken ? resp?.nextToken : ''
             sessionId = resp?.$response?.httpResponse?.headers['x-amzn-sessionid']
+            getLogger().debug(`${isNextSession ? 'current' : 'next'} session id: ${sessionId}`)
+
             TelemetryHelper.instance.setFirstResponseRequestId(requestId)
             if (page === 0) {
                 currentSession.setTimeToFirstRecommendation(performance.now())
@@ -390,6 +394,7 @@ export class RecommendationHandler {
 
         if (!isNextSession && this.isCancellationRequested()) {
             return Promise.resolve<GetRecommendationsResponse>({
+                sessionId: '',
                 result: invocationResult,
                 errorMessage: errorMessage,
                 recommendationCount: currentSession.recommendations.length,
@@ -450,6 +455,7 @@ export class RecommendationHandler {
             }
         }
         return Promise.resolve<GetRecommendationsResponse>({
+            sessionId: sessionId,
             result: invocationResult,
             errorMessage: errorMessage,
             recommendationCount: currentSession.recommendations.length,
@@ -638,7 +644,7 @@ export class RecommendationHandler {
         const session = CodeWhispererSessionState.instance.getSession()
 
         if (!indexShift && session.recommendations.length) {
-            await this.fetchNextRecommendations()
+            await this.fetchNextRecommendations(session)
         }
         await lock.acquire(updateInlineLockKey, async () => {
             if (!vscode.window.state.focused) {
@@ -717,15 +723,14 @@ export class RecommendationHandler {
         }
     }
 
-    async fetchNextRecommendations() {
-        const session = CodeWhispererSessionState.instance.getSession()
+    async fetchNextRecommendations(session: CodeWhispererSession) {
         const client = new codewhispererClient.DefaultCodeWhispererClient()
         const editor = vscode.window.activeTextEditor
         if (!editor) {
             return
         }
 
-        await this.getRecommendations(
+        const nextResp = await this.getRecommendations(
             client,
             editor,
             session.triggerType,
@@ -737,6 +742,14 @@ export class RecommendationHandler {
             false,
             true
         )
+
+        const nextSessionId = nextResp.sessionId
+        if (nextSessionId.length === 0 || nextSessionId !== session.sessionId) {
+            getLogger().debug(
+                `next session id ${nextSessionId} mismatch previous session id ${session.sessionId}, resetting next session`
+            )
+            CodeWhispererSessionState.instance.getNextSession().reset()
+        }
     }
 
     async tryShowRecommendation() {